Annotated the return type of 'elvis' function with @Exact

fun <T> ELVIS(T?, T): @Exact T

compiler/frontend/src/org/jetbrains/kotlin/descriptors/annotations/AnnotationSplitter.kt

@@ -127,6 +127,7 @@ public class AnnotationSplitter(
         }
 
         override fun isEmpty() = annotations.isEmpty()
+        override fun hasAnnotation(fqName: FqName) = annotations.hasAnnotation(fqName)
         override fun findAnnotation(fqName: FqName) = annotations.findAnnotation(fqName)
         override fun getUseSiteTargetedAnnotations() = annotations.getUseSiteTargetedAnnotations()
         override fun getAllAnnotations() = annotations.getAllAnnotations()

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

@@ -63,6 +63,7 @@ import org.jetbrains.kotlin.resolve.scopes.LexicalWritableScope;
 import org.jetbrains.kotlin.resolve.scopes.receivers.ExpressionReceiver;
 import org.jetbrains.kotlin.resolve.scopes.utils.ScopeUtilsKt;
 import org.jetbrains.kotlin.types.*;
+import org.jetbrains.kotlin.types.expressions.ControlStructureTypingUtils.ResolveConstruct;
 import org.jetbrains.kotlin.types.checker.KotlinTypeChecker;
 import org.jetbrains.kotlin.types.expressions.typeInfoFactory.TypeInfoFactoryKt;
 import org.jetbrains.kotlin.types.expressions.unqualifiedSuper.UnqualifiedSuperKt;
@@ -830,7 +831,7 @@ public class BasicExpressionTypingVisitor extends ExpressionTypingVisitor {
         // See also CallExpressionResolver.getSimpleNameExpressionTypeInfo, .getQualifiedExpressionTypeInfo
         Call call = createCallForSpecialConstruction(expression, expression.getOperationReference(), Collections.singletonList(baseExpression));
         components.controlStructureTypingUtils.resolveSpecialConstructionAsCall(
-                call, "ExclExcl", Collections.singletonList("baseExpr"), Collections.singletonList(true), context, null);
+                call, ResolveConstruct.EXCL_EXCL, Collections.singletonList("baseExpr"), Collections.singletonList(true), context, null);
         JetTypeInfo baseTypeInfo = BindingContextUtils.getRecordedTypeInfo(baseExpression, context.trace.getBindingContext());
 
         if (ArgumentTypeResolver.isFunctionLiteralArgument(baseExpression, context)) {
@@ -1163,7 +1164,7 @@ public class BasicExpressionTypingVisitor extends ExpressionTypingVisitor {
 
         Call call = createCallForSpecialConstruction(expression, expression.getOperationReference(), Lists.newArrayList(left, right));
         ResolvedCall<FunctionDescriptor> resolvedCall = components.controlStructureTypingUtils.resolveSpecialConstructionAsCall(
-                call, "Elvis", Lists.newArrayList("left", "right"), Lists.newArrayList(true, false), contextWithExpectedType, null);
+                call, ResolveConstruct.ELVIS, Lists.newArrayList("left", "right"), Lists.newArrayList(true, false), contextWithExpectedType, null);
         JetTypeInfo leftTypeInfo = BindingContextUtils.getRecordedTypeInfo(left, context.trace.getBindingContext());
         if (ArgumentTypeResolver.isFunctionLiteralArgument(left, context)) {
             context.trace.report(USELESS_ELVIS_ON_FUNCTION_LITERAL.on(expression.getOperationReference()));

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

@@ -47,10 +47,12 @@ import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind;
 import org.jetbrains.kotlin.resolve.calls.tasks.ResolutionCandidate;
 import org.jetbrains.kotlin.resolve.calls.tasks.TracingStrategy;
 import org.jetbrains.kotlin.resolve.calls.util.CallMaker;
+import org.jetbrains.kotlin.resolve.descriptorUtil.AnnotationsForResolveKt;
 import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue;
 import org.jetbrains.kotlin.types.KotlinType;
 import org.jetbrains.kotlin.types.TypeUtils;
 import org.jetbrains.kotlin.types.Variance;
+import org.jetbrains.kotlin.types.typeUtil.TypeUtilsKt;
 
 import java.util.*;
 
@@ -61,6 +63,20 @@ import static org.jetbrains.kotlin.resolve.calls.inference.constraintPosition.Co
 public class ControlStructureTypingUtils {
     private static final Logger LOG = Logger.getInstance(ControlStructureTypingUtils.class);
 
+    public enum ResolveConstruct {
+        IF("if"), ELVIS("elvis"), EXCL_EXCL("ExclExcl");
+
+        private final String name;
+
+        ResolveConstruct(String name) {
+            this.name = name;
+        }
+
+        public String getName() {
+            return name;
+        }
+    }
+
     private final CallResolver callResolver;
     private final DataFlowAnalyzer dataFlowAnalyzer;
     private final ModuleDescriptor moduleDescriptor;
@@ -77,15 +93,15 @@ public class ControlStructureTypingUtils {
 
     /*package*/ ResolvedCall<FunctionDescriptor> resolveSpecialConstructionAsCall(
             @NotNull Call call,
-            @NotNull String constructionName,
+            @NotNull ResolveConstruct construct,
             @NotNull List<String> argumentNames,
             @NotNull List<Boolean> isArgumentNullable,
             @NotNull ExpressionTypingContext context,
             @Nullable MutableDataFlowInfoForArguments dataFlowInfoForArguments
     ) {
         SimpleFunctionDescriptorImpl function = createFunctionDescriptorForSpecialConstruction(
-                constructionName.toUpperCase(), argumentNames, isArgumentNullable);
+                construct, argumentNames, isArgumentNullable);
-        TracingStrategy tracing = createTracingForSpecialConstruction(call, constructionName, context);
+        TracingStrategy tracing = createTracingForSpecialConstruction(call, construct.getName(), context);
         ResolutionCandidate<CallableDescriptor> resolutionCandidate = ResolutionCandidate.<CallableDescriptor>create(call, function);
         OverloadResolutionResults<FunctionDescriptor> results = callResolver.resolveCallWithKnownCandidate(
                 call, tracing, context, resolutionCandidate, dataFlowInfoForArguments);
@@ -94,12 +110,13 @@ public class ControlStructureTypingUtils {
     }
 
     private SimpleFunctionDescriptorImpl createFunctionDescriptorForSpecialConstruction(
-            @NotNull String constructionName,
+            @NotNull ResolveConstruct construct,
             @NotNull List<String> argumentNames,
             @NotNull List<Boolean> isArgumentNullable
     ) {
         assert argumentNames.size() == isArgumentNullable.size();
 
+        String constructionName = construct.getName().toUpperCase();
         Name specialFunctionName = Name.identifier("<SPECIAL-FUNCTION-FOR-" + constructionName + "-RESOLVE>");
 
         SimpleFunctionDescriptorImpl function = SimpleFunctionDescriptorImpl.create(
@@ -126,12 +143,13 @@ public class ControlStructureTypingUtils {
             );
             valueParameters.add(valueParameter);
         }
+        KotlinType returnType = construct != ResolveConstruct.ELVIS ? type : TypeUtilsKt.replaceAnnotations(type, AnnotationsForResolveKt.getExactInAnnotations());
         function.initialize(
                 null,
                 null,
                 Lists.newArrayList(typeParameter),
                 valueParameters,
-                type,
+                returnType,
                 Modality.FINAL,
                 Visibilities.PUBLIC
         );

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

@@ -41,6 +41,7 @@ import org.jetbrains.kotlin.types.CommonSupertypes;
 import org.jetbrains.kotlin.types.ErrorUtils;
 import org.jetbrains.kotlin.types.KotlinType;
 import org.jetbrains.kotlin.types.TypeUtils;
+import org.jetbrains.kotlin.types.expressions.ControlStructureTypingUtils.ResolveConstruct;
 import org.jetbrains.kotlin.types.checker.KotlinTypeChecker;
 import org.jetbrains.kotlin.types.expressions.typeInfoFactory.TypeInfoFactoryKt;
 
@@ -128,7 +129,7 @@ public class ControlStructureTypingVisitor extends ExpressionTypingVisitor {
         MutableDataFlowInfoForArguments dataFlowInfoForArguments =
                     createDataFlowInfoForArgumentsForIfCall(callForIf, thenInfo, elseInfo);
         ResolvedCall<FunctionDescriptor> resolvedCall = components.controlStructureTypingUtils.resolveSpecialConstructionAsCall(
-                callForIf, "If", Lists.newArrayList("thenBranch", "elseBranch"),
+                callForIf, ResolveConstruct.IF, Lists.newArrayList("thenBranch", "elseBranch"),
                 Lists.newArrayList(false, false),
                 contextWithExpectedType, dataFlowInfoForArguments);
 

compiler/testData/diagnostics/tests/controlStructures/lambdasInExclExclAndElvis.kt

@@ -8,8 +8,8 @@ fun test() {
 
     // KT-KT-9070
     <!TYPE_MISMATCH!>{ }<!> <!USELESS_ELVIS!><!USELESS_ELVIS_ON_FUNCTION_LITERAL!>?:<!> 1<!>
-    use(<!TYPE_MISMATCH!>{ 2 }<!> <!USELESS_ELVIS_ON_FUNCTION_LITERAL!>?:<!> 1);
+    use({ 2 } <!USELESS_ELVIS_ON_FUNCTION_LITERAL!>?:<!> 1);
 
     1 <!USELESS_ELVIS!>?: <!TYPE_MISMATCH, UNUSED_FUNCTION_LITERAL!>{ }<!><!>
-    use(1 <!USELESS_ELVIS!>?: <!TYPE_MISMATCH!>{ }<!><!>)
+    use(1 <!USELESS_ELVIS!>?: { }<!>)
 }

compiler/testData/diagnostics/tests/resolve/specialConstructions/inferenceForElvis.kt

@@ -0,0 +1,8 @@
+// !DIAGNOSTICS: -UNUSED_PARAMETER -UNUSED_VARIABLE
+// !CHECK_TYPE
+
+fun <T: Any> foo(f: (T) -> Unit): T? = null // T is used only as return type
+fun test() {
+    val x = foo { it checkType { _<String>() }} ?: "" // foo() is inferred as foo<String>, which isn't very good
+    val y: Any = foo { it checkType { _<Any>() } } ?: "" // but for now it's fixed by specifying expected type
+}

compiler/testData/diagnostics/tests/resolve/specialConstructions/inferenceForElvis.txt

@@ -0,0 +1,4 @@
+package
+
+public fun </*0*/ T : kotlin.Any> foo(/*0*/ f: (T) -> kotlin.Unit): T?
+public fun test(): kotlin.Unit

compiler/tests/org/jetbrains/kotlin/checkers/JetDiagnosticsTestGenerated.java

@@ -13088,6 +13088,12 @@ public class JetDiagnosticsTestGenerated extends AbstractJetDiagnosticsTest {
                     doTest(fileName);
                 }
 
+                @TestMetadata("inferenceForElvis.kt")
+                public void testInferenceForElvis() throws Exception {
+                    String fileName = JetTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/resolve/specialConstructions/inferenceForElvis.kt");
+                    doTest(fileName);
+                }
+
                 @TestMetadata("multipleSuperClasses.kt")
                 public void testMultipleSuperClasses() throws Exception {
                     String fileName = JetTestUtils.navigationMetadata("compiler/testData/diagnostics/tests/resolve/specialConstructions/multipleSuperClasses.kt");

core/descriptors/src/org/jetbrains/kotlin/descriptors/annotations/Annotations.kt

@@ -80,6 +80,11 @@ class FilteredAnnotations(
         private val delegate: Annotations,
         private val fqNameFilter: (FqName) -> Boolean
 ) : Annotations {
+
+    override fun hasAnnotation(fqName: FqName) =
+            if (fqNameFilter(fqName)) delegate.hasAnnotation(fqName)
+            else false
+
     override fun findAnnotation(fqName: FqName) =
             if (fqNameFilter(fqName)) delegate.findAnnotation(fqName)
             else null
@@ -115,6 +120,8 @@ class CompositeAnnotations(
 
     override fun isEmpty() = delegates.all { it.isEmpty() }
 
+    override fun hasAnnotation(fqName: FqName) = delegates.asSequence().any { it.hasAnnotation(fqName) }
+
     override fun findAnnotation(fqName: FqName) = delegates.asSequence().map { it.findAnnotation(fqName) }.filterNotNull().firstOrNull()
 
     override fun findExternalAnnotation(fqName: FqName) = delegates.asSequence().map { it.findExternalAnnotation(fqName) }.filterNotNull().firstOrNull()

core/descriptors/src/org/jetbrains/kotlin/resolve/annotationsForResolve.kt

@@ -18,6 +18,8 @@ package org.jetbrains.kotlin.resolve.descriptorUtil
 
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.descriptors.TypeParameterDescriptor
+import org.jetbrains.kotlin.descriptors.annotations.AnnotationDescriptor
+import org.jetbrains.kotlin.descriptors.annotations.AnnotationWithTarget
 import org.jetbrains.kotlin.descriptors.annotations.Annotations
 import org.jetbrains.kotlin.name.FqName
 import org.jetbrains.kotlin.types.KotlinType
@@ -41,3 +43,21 @@ public fun CallableDescriptor.hasLowPriorityInOverloadResolution(): Boolean = an
 private val ONLY_INPUT_TYPES_FQ_NAME = FqName("kotlin.internal.OnlyInputTypes")
 
 public fun TypeParameterDescriptor.hasOnlyInputTypesAnnotation(): Boolean = annotations.hasAnnotation(ONLY_INPUT_TYPES_FQ_NAME)
+
+public fun getExactInAnnotations(): Annotations = AnnotationsWithOnly(EXACT_ANNOTATION_FQ_NAME)
+
+private class AnnotationsWithOnly(val presentAnnotation: FqName): Annotations {
+    override fun iterator(): Iterator<AnnotationDescriptor> = listOf<AnnotationDescriptor>().iterator()
+
+    override fun isEmpty(): Boolean = false
+
+    override fun hasAnnotation(fqName: FqName): Boolean = fqName == this.presentAnnotation
+
+    override fun findAnnotation(fqName: FqName): AnnotationDescriptor? = null
+
+    override fun findExternalAnnotation(fqName: FqName): AnnotationDescriptor? = null
+
+    override fun getUseSiteTargetedAnnotations(): List<AnnotationWithTarget> = emptyList()
+
+    override fun getAllAnnotations(): List<AnnotationWithTarget> = emptyList()
+}

js/js.translator/testData/expression/misc/cases/elvisReturnSimple.kt

@@ -12,7 +12,7 @@ fun stringLen(s : String?) : Int {
 }
 
 fun stringReturnInLeftLen(s : String?) : Int {
-    val s1 : String = (if (s != null) { return s.length() } else { null }) ?: return 0
+    val s1 = (if (s != null) { return s.length() } else { null }) ?: return 0
 }
 
 fun box(): String {