Emit not-null assertions for enhanced types

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

@@ -37,6 +37,7 @@ import org.jetbrains.kotlin.load.java.JavaVisibilities;
 import org.jetbrains.kotlin.load.java.JvmAbi;
 import org.jetbrains.kotlin.load.java.JvmAnnotationNames;
 import org.jetbrains.kotlin.load.java.descriptors.JavaCallableMemberDescriptor;
+import org.jetbrains.kotlin.jvm.RuntimeAssertionInfo;
 import org.jetbrains.kotlin.name.FqName;
 import org.jetbrains.kotlin.resolve.DescriptorUtils;
 import org.jetbrains.kotlin.resolve.annotations.AnnotationsPackage;
@@ -45,7 +46,6 @@ import org.jetbrains.kotlin.resolve.jvm.JvmClassName;
 import org.jetbrains.kotlin.resolve.jvm.JvmPackage;
 import org.jetbrains.kotlin.resolve.jvm.JvmPrimitiveType;
 import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmDeclarationOrigin;
-import org.jetbrains.kotlin.types.Approximation;
 import org.jetbrains.kotlin.types.JetType;
 import org.jetbrains.kotlin.types.TypesPackage;
 import org.jetbrains.org.objectweb.asm.*;
@@ -637,10 +637,10 @@ public class AsmUtil {
     public static StackValue genNotNullAssertions(
             @NotNull GenerationState state,
             @NotNull final StackValue stackValue,
-            @Nullable final Approximation.Info approximationInfo
+            @Nullable final RuntimeAssertionInfo runtimeAssertionInfo
     ) {
         if (!state.isCallAssertionsEnabled()) return stackValue;
-        if (approximationInfo == null || !TypesPackage.assertNotNull(approximationInfo)) return stackValue;
+        if (runtimeAssertionInfo == null || !runtimeAssertionInfo.getNeedNotNullAssertion()) return stackValue;
 
         return new StackValue(stackValue.type) {
 
@@ -649,7 +649,7 @@ public class AsmUtil {
                 stackValue.put(type, v);
                 if (type.getSort() == Type.OBJECT || type.getSort() == Type.ARRAY) {
                     v.dup();
-                    v.visitLdcInsn(approximationInfo.getMessage());
+                    v.visitLdcInsn(runtimeAssertionInfo.getMessage());
                     v.invokestatic("kotlin/jvm/internal/Intrinsics", "checkExpressionValueIsNotNull",
                                    "(Ljava/lang/Object;Ljava/lang/String;)V", false);
                 }

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

@@ -49,10 +49,12 @@ import org.jetbrains.kotlin.descriptors.*;
 import org.jetbrains.kotlin.descriptors.impl.ScriptCodeDescriptor;
 import org.jetbrains.kotlin.diagnostics.DiagnosticUtils;
 import org.jetbrains.kotlin.diagnostics.Errors;
+import org.jetbrains.kotlin.jvm.bindingContextSlices.BindingContextSlicesPackage;
 import org.jetbrains.kotlin.lexer.JetTokens;
 import org.jetbrains.kotlin.load.java.JvmAbi;
 import org.jetbrains.kotlin.load.java.descriptors.JavaClassDescriptor;
 import org.jetbrains.kotlin.load.java.descriptors.SamConstructorDescriptor;
+import org.jetbrains.kotlin.jvm.RuntimeAssertionInfo;
 import org.jetbrains.kotlin.name.Name;
 import org.jetbrains.kotlin.psi.*;
 import org.jetbrains.kotlin.renderer.DescriptorRenderer;
@@ -75,7 +77,6 @@ import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodParameterKind;
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodParameterSignature;
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodSignature;
 import org.jetbrains.kotlin.resolve.scopes.receivers.*;
-import org.jetbrains.kotlin.types.Approximation;
 import org.jetbrains.kotlin.types.JetType;
 import org.jetbrains.kotlin.types.TypeProjection;
 import org.jetbrains.kotlin.types.TypeUtils;
@@ -280,12 +281,12 @@ public class ExpressionCodegen extends JetVisitor<StackValue, StackValue> implem
 
             StackValue stackValue = selector.accept(visitor, receiver);
 
-            Approximation.Info approximationInfo = null;
+            RuntimeAssertionInfo runtimeAssertionInfo = null;
             if (selector instanceof JetExpression) {
-                approximationInfo = bindingContext.get(BindingContext.EXPRESSION_RESULT_APPROXIMATION, (JetExpression) selector);
+                runtimeAssertionInfo = bindingContext.get(BindingContextSlicesPackage.getRUNTIME_ASSERTION_INFO(), (JetExpression) selector);
             }
 
-            return genNotNullAssertions(state, stackValue, approximationInfo);
+            return genNotNullAssertions(state, stackValue, runtimeAssertionInfo);
         }
         catch (ProcessCanceledException e) {
             throw e;

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

@@ -35,6 +35,7 @@ import org.jetbrains.kotlin.codegen.state.GenerationState;
 import org.jetbrains.kotlin.codegen.state.JetTypeMapper;
 import org.jetbrains.kotlin.descriptors.*;
 import org.jetbrains.kotlin.descriptors.annotations.AnnotationDescriptor;
+import org.jetbrains.kotlin.jvm.RuntimeAssertionInfo;
 import org.jetbrains.kotlin.load.kotlin.nativeDeclarations.NativeDeclarationsPackage;
 import org.jetbrains.kotlin.name.FqName;
 import org.jetbrains.kotlin.psi.JetNamedFunction;
@@ -51,8 +52,6 @@ import org.jetbrains.kotlin.resolve.jvm.diagnostics.JvmDeclarationOrigin;
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodParameterKind;
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodParameterSignature;
 import org.jetbrains.kotlin.resolve.jvm.jvmSignature.JvmMethodSignature;
-import org.jetbrains.kotlin.types.Approximation;
-import org.jetbrains.kotlin.types.TypesPackage;
 import org.jetbrains.org.objectweb.asm.AnnotationVisitor;
 import org.jetbrains.org.objectweb.asm.Label;
 import org.jetbrains.org.objectweb.asm.MethodVisitor;
@@ -834,10 +833,10 @@ public class FunctionCodegen {
                         StackValue stackValue = AsmUtil.genNotNullAssertions(
                                 state,
                                 StackValue.onStack(delegateToMethod.getReturnType()),
-                                TypesPackage.getApproximationTo(
-                                        delegatedTo.getReturnType(),
+                                RuntimeAssertionInfo.create(
                                         delegateFunction.getReturnType(),
-                                        new Approximation.DataFlowExtras.OnlyMessage(delegatedTo.getName() + "(...)")
+                                        delegatedTo.getReturnType(),
+                                        new RuntimeAssertionInfo.DataFlowExtras.OnlyMessage(delegatedTo.getName() + "(...)")
                                 )
                         );
 

compiler/frontend.java/src/org/jetbrains/kotlin/jvm/RuntimeAssertions.kt

@@ -0,0 +1,116 @@
+/*
+ * Copyright 2010-2015 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.jvm
+
+import com.intellij.openapi.util.text.StringUtil
+import org.jetbrains.kotlin.descriptors.ReceiverParameterDescriptor
+import org.jetbrains.kotlin.load.java.JvmAnnotationNames
+import org.jetbrains.kotlin.jvm.bindingContextSlices.RUNTIME_ASSERTION_INFO
+import org.jetbrains.kotlin.psi.JetExpression
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.calls.checkers.AdditionalTypeChecker
+import org.jetbrains.kotlin.resolve.calls.context.CallResolutionContext
+import org.jetbrains.kotlin.resolve.calls.context.ResolutionContext
+import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValueFactory
+import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue
+import org.jetbrains.kotlin.types.JetType
+import org.jetbrains.kotlin.types.TypeUtils
+import org.jetbrains.kotlin.types.checker.JetTypeChecker
+import org.jetbrains.kotlin.types.upperIfFlexible
+import kotlin.platform.platformStatic
+
+public class RuntimeAssertionInfo(public val needNotNullAssertion: Boolean, public val message: String) {
+    public interface DataFlowExtras {
+        class OnlyMessage(message: String) : DataFlowExtras {
+            override val canBeNull: Boolean get() = true
+            override val possibleTypes: Set<JetType> get() = setOf()
+            override val presentableText: String = message
+        }
+
+        val canBeNull: Boolean
+        val possibleTypes: Set<JetType>
+        val presentableText: String
+    }
+
+    companion object {
+        platformStatic public fun create(
+                expectedType: JetType,
+                expressionType: JetType,
+                dataFlowExtras: DataFlowExtras
+        ): RuntimeAssertionInfo? {
+            fun assertNotNull(): Boolean {
+                if (expectedType.isError() || expressionType.isError()) return false
+
+                // T : Any, T! = T..T?
+                // Let T$ will be copy of T! with enhanced nullability.
+                // Cases when nullability assertion needed: T! -> T, T$ -> T
+
+                // Expected type either T?, T! or T$
+                if (TypeUtils.isNullableType(expectedType) || expectedType.hasEnhancedNullability()) return false
+
+                // Expression type is not nullable and not enhanced (neither T?, T! or T$)
+                val isExpressionTypeNullable = TypeUtils.isNullableType(expressionType)
+                if (!isExpressionTypeNullable && !expressionType.hasEnhancedNullability()) return false
+
+                // Smart-cast T! or T?
+                if (!dataFlowExtras.canBeNull && isExpressionTypeNullable) return false
+
+                return true
+            }
+
+            return if (assertNotNull())
+                RuntimeAssertionInfo(needNotNullAssertion = true, message = dataFlowExtras.presentableText)
+            else
+                null
+        }
+
+        private fun JetType.hasEnhancedNullability()
+                = getAnnotations().findAnnotation(JvmAnnotationNames.ENHANCED_NULLABILITY_ANNOTATION) != null
+    }
+}
+
+public object RuntimeAssertionsTypeChecker : AdditionalTypeChecker {
+    override fun checkType(expression: JetExpression, expressionType: JetType, c: ResolutionContext<*>) {
+        if (TypeUtils.noExpectedType(c.expectedType)) return
+
+        val assertionInfo = RuntimeAssertionInfo.create(
+                c.expectedType,
+                expressionType,
+                object : RuntimeAssertionInfo.DataFlowExtras {
+                    override val canBeNull: Boolean
+                        get() = c.dataFlowInfo.getNullability(dataFlowValue).canBeNull()
+                    override val possibleTypes: Set<JetType>
+                        get() = c.dataFlowInfo.getPossibleTypes(dataFlowValue)
+                    override val presentableText: String
+                        get() = StringUtil.trimMiddle(expression.getText(), 50)
+
+                    private val dataFlowValue = DataFlowValueFactory.createDataFlowValue(expression, expressionType, c)
+                }
+        )
+
+        if (assertionInfo != null) {
+            c.trace.record(RUNTIME_ASSERTION_INFO, expression, assertionInfo)
+        }
+    }
+
+    override fun checkReceiver(
+            receiverParameter: ReceiverParameterDescriptor,
+            receiverArgument: ReceiverValue,
+            safeAccess: Boolean,
+            c: CallResolutionContext<*>
+    ) { }
+}

compiler/frontend.java/src/org/jetbrains/kotlin/jvm/bindingContextSlices.kt

@@ -0,0 +1,26 @@
+/*
+ * Copyright 2010-2015 JetBrains s.r.o.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.jetbrains.kotlin.jvm.bindingContextSlices
+
+import org.jetbrains.kotlin.jvm.RuntimeAssertionInfo
+import org.jetbrains.kotlin.psi.JetExpression
+import org.jetbrains.kotlin.util.slicedMap.BasicWritableSlice
+import org.jetbrains.kotlin.util.slicedMap.RewritePolicy
+import org.jetbrains.kotlin.util.slicedMap.WritableSlice
+import org.jetbrains.kotlin.utils.DO_NOTHING
+
+public val RUNTIME_ASSERTION_INFO: WritableSlice<JetExpression, RuntimeAssertionInfo> = BasicWritableSlice(RewritePolicy.DO_NOTHING)

compiler/frontend.java/src/org/jetbrains/kotlin/load/kotlin/KotlinJvmCheckerProvider.kt

@@ -20,6 +20,7 @@ import org.jetbrains.kotlin.cfg.WhenChecker
 import org.jetbrains.kotlin.descriptors.*
 import org.jetbrains.kotlin.diagnostics.DiagnosticSink
 import org.jetbrains.kotlin.diagnostics.Errors
+import org.jetbrains.kotlin.jvm.RuntimeAssertionsTypeChecker
 import org.jetbrains.kotlin.lexer.JetTokens
 import org.jetbrains.kotlin.load.java.lazy.types.isMarkedNotNull
 import org.jetbrains.kotlin.load.java.lazy.types.isMarkedNullable
@@ -63,7 +64,7 @@ public class KotlinJvmCheckerProvider(private val module: ModuleDescriptor) : Ad
                                         JavaAnnotationMethodCallChecker(), TraitDefaultMethodCallChecker(),
                                         ReflectionAPICallChecker(module)),
 
-        additionalTypeCheckers = listOf(JavaNullabilityWarningsChecker()),
+        additionalTypeCheckers = listOf(JavaNullabilityWarningsChecker(), RuntimeAssertionsTypeChecker),
         additionalSymbolUsageValidators = listOf()
 )
 

compiler/testData/codegen/notNullAssertions/AssertionChecker.kt

@@ -45,10 +45,10 @@ fun checkAssertions(illegalStateExpected: Boolean) {
     check("plus") { A() + A() }
     
     // field
-    // check("NULL") { A().NULL }
+    check("NULL") { A().NULL }
     
     // static field
-    // check("STATIC_NULL") { A.STATIC_NULL }
+    check("STATIC_NULL") { A.STATIC_NULL }
 
     // postfix expression
     // TODO:

compiler/tests/org/jetbrains/kotlin/codegen/GenerateNotNullAssertionsTest.java

@@ -126,6 +126,7 @@ public class GenerateNotNullAssertionsTest extends CodegenTestCase {
         loadFile("notNullAssertions/arrayListGet.kt");
         String text = generateToText();
 
+        assertTrue(text.contains("checkExpressionValueIsNotNull"));
         assertTrue(text.contains("checkParameterIsNotNull"));
     }
 
@@ -136,7 +137,7 @@ public class GenerateNotNullAssertionsTest extends CodegenTestCase {
         loadFile("notNullAssertions/javaMultipleSubstitutions.kt");
         String text = generateToText();
 
-        assertEquals(0, StringUtil.getOccurrenceCount(text, "checkExpressionValueIsNotNull"));
+        assertEquals(3, StringUtil.getOccurrenceCount(text, "checkExpressionValueIsNotNull"));
         assertEquals(3, StringUtil.getOccurrenceCount(text, "checkParameterIsNotNull"));
     }