Move detecting of illegal call of contract function into call checker

(#KT-26153, #KT-26191) fixed

compiler/frontend/src/org/jetbrains/kotlin/contracts/parsing/ContractParsingDiagnosticsCollector.kt

@@ -37,6 +37,8 @@ interface ContractParsingDiagnosticsCollector {
 }
 
 class TraceBasedCollector(private val bindingTrace: BindingTrace, mainCall: KtExpression) : ContractParsingDiagnosticsCollector {
+    constructor(callContext: ContractCallContext) : this(callContext.trace, callContext.contractCallExpression)
+
     private val reportedErrors: MutableList<Diagnostic> = mutableListOf()
     private val mainCallReportTarget = (mainCall as? KtCallExpression)?.calleeExpression ?: mainCall
 

compiler/frontend/src/org/jetbrains/kotlin/contracts/parsing/ContractParsingServices.kt

@@ -23,13 +23,12 @@ import org.jetbrains.kotlin.contracts.description.ContractDescription
 import org.jetbrains.kotlin.contracts.description.ContractProviderKey
 import org.jetbrains.kotlin.contracts.description.LazyContractProvider
 import org.jetbrains.kotlin.descriptors.*
-import org.jetbrains.kotlin.psi.KtExpression
+import org.jetbrains.kotlin.psi.*
 import org.jetbrains.kotlin.psi.psiUtil.isContractDescriptionCallPsiCheck
 import org.jetbrains.kotlin.resolve.BindingContext
 import org.jetbrains.kotlin.resolve.BindingTrace
 import org.jetbrains.kotlin.resolve.calls.callUtil.getResolvedCall
 import org.jetbrains.kotlin.resolve.scopes.LexicalScope
-import org.jetbrains.kotlin.resolve.scopes.LexicalScopeKind
 import org.jetbrains.kotlin.storage.StorageManager
 
 class ContractParsingServices(val languageVersionSettings: LanguageVersionSettings, private val storageManager: StorageManager) {
@@ -42,12 +41,12 @@ class ContractParsingServices(val languageVersionSettings: LanguageVersionSettin
      *
      * Otherwise, it may lead to inconsistent resolve state and failed assertions
      */
-    fun checkContractAndRecordIfPresent(expression: KtExpression, trace: BindingTrace, scope: LexicalScope, isFirstStatement: Boolean) {
+    fun checkContractAndRecordIfPresent(expression: KtExpression, trace: BindingTrace, scope: LexicalScope) {
         // Fastpath. Note that it doesn't violates invariant described in KDoc, because 'isContractDescriptionCallPsiCheck'
         // is a *necessary* (but not sufficient, actually) condition for presence of 'LazyContractProvider'
         if (!expression.isContractDescriptionCallPsiCheck()) return
 
-        val callContext = ContractCallContext(expression, isFirstStatement, scope, trace)
+        val callContext = ContractCallContext(expression, scope, trace)
         val contractProviderIfAny =
             (scope.ownerDescriptor as? FunctionDescriptor)?.getUserData(ContractProviderKey) as? LazyContractProvider?
         var resultingContractDescription: ContractDescription? = null
@@ -68,14 +67,14 @@ class ContractParsingServices(val languageVersionSettings: LanguageVersionSettin
      * ideally, it should satisfy following condition: null returned <=> at least one error was reported
      */
     private fun parseContractAndReportErrors(callContext: ContractCallContext): ContractDescription? {
-        val collector = TraceBasedCollector(callContext.trace, callContext.contractCallExpression)
+        val collector = TraceBasedCollector(callContext)
 
         try {
             checkFeatureEnabled(collector)
-            checkContractAllowedHere(collector, callContext)
 
+            val contractNotAllowed = callContext.bindingContext[BindingContext.CONTRACT_NOT_ALLOWED, callContext.contractCallExpression] == true
             // Small optimization: do not even try to parse contract if we already have errors
-            if (collector.hasErrors()) return null
+            if (collector.hasErrors() || contractNotAllowed) return null
 
             val parsedContract = PsiContractParserDispatcher(collector, callContext, storageManager).parseContract()
 
@@ -100,35 +99,12 @@ class ContractParsingServices(val languageVersionSettings: LanguageVersionSettin
         }
     }
 
-    private fun checkContractAllowedHere(collector: ContractParsingDiagnosticsCollector, callContext: ContractCallContext) {
-        val functionDescriptor = callContext.ownerDescriptor as? FunctionDescriptor
-        val scope = callContext.scope
-
-        if (!callContext.isFirstStatement)
-            collector.contractNotAllowed("Contract should be the first statement")
-
-        if (functionDescriptor == null || functionDescriptor is PropertyAccessorDescriptor)
-            collector.contractNotAllowed("Contracts are allowed only for functions")
-
-
-        if (callContext.ownerDescriptor.containingDeclaration !is PackageFragmentDescriptor
-            || scope.kind != LexicalScopeKind.CODE_BLOCK
-            || (scope.parent as? LexicalScope)?.kind != LexicalScopeKind.FUNCTION_INNER_SCOPE
-        )
-            collector.contractNotAllowed("Contracts are allowed only for top-level functions")
-
-        if (functionDescriptor?.isOperator == true) collector.contractNotAllowed("Contracts are not allowed for operator functions")
-
-        if (functionDescriptor?.isOverridable == true) collector.contractNotAllowed("Contracts are not allowed for open functions")
-    }
-
     private fun KtExpression.isContractDescriptionCallPreciseCheck(context: BindingContext): Boolean =
         getResolvedCall(context)?.resultingDescriptor?.isContractCallDescriptor() ?: false
 }
 
 class ContractCallContext(
     val contractCallExpression: KtExpression,
-    val isFirstStatement: Boolean,
     val scope: LexicalScope,
     val trace: BindingTrace
 ) {

compiler/frontend/src/org/jetbrains/kotlin/resolve/BindingContext.java

@@ -98,6 +98,7 @@ public interface BindingContext {
     WritableSlice<KtExpression, DataFlowInfo> DATA_FLOW_INFO_BEFORE = new BasicWritableSlice<>(DO_NOTHING);
     WritableSlice<KtExpression, KotlinType> EXPECTED_EXPRESSION_TYPE = new BasicWritableSlice<>(DO_NOTHING);
     WritableSlice<KtElement, Computation> EXPRESSION_EFFECTS = Slices.createSimpleSlice();
+    WritableSlice<KtElement, Boolean> CONTRACT_NOT_ALLOWED = Slices.createSimpleSlice();
     WritableSlice<KtFunction, KotlinType> EXPECTED_RETURN_TYPE = new BasicWritableSlice<>(DO_NOTHING);
     WritableSlice<KtExpression, DataFlowInfo> DATAFLOW_INFO_AFTER_CONDITION = Slices.createSimpleSlice();
     WritableSlice<VariableDescriptor, DataFlowValue> BOUND_INITIALIZER_VALUE = Slices.createSimpleSlice();

compiler/frontend/src/org/jetbrains/kotlin/resolve/PlatformConfiguratorBase.kt

@@ -46,7 +46,7 @@ private val DEFAULT_CALL_CHECKERS = listOf(
     UnderscoreUsageChecker, AssigningNamedArgumentToVarargChecker(),
     PrimitiveNumericComparisonCallChecker, LambdaWithSuspendModifierCallChecker,
     UselessElvisCallChecker(), ResultTypeWithNullableOperatorsChecker(), NullableVarargArgumentCallChecker,
-    NamedFunAsExpressionChecker
+    NamedFunAsExpressionChecker, ContractNotAllowedCallChecker
 )
 private val DEFAULT_TYPE_CHECKERS = emptyList<AdditionalTypeChecker>()
 private val DEFAULT_CLASSIFIER_USAGE_CHECKERS = listOf(

compiler/frontend/src/org/jetbrains/kotlin/resolve/calls/checkers/ContractNotAllowedCallChecker.kt

@@ -0,0 +1,74 @@
+/*
+ * Copyright 2010-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license
+ * that can be found in the license/LICENSE.txt file.
+ */
+
+package org.jetbrains.kotlin.resolve.calls.checkers
+
+import com.intellij.psi.PsiElement
+import org.jetbrains.kotlin.contracts.parsing.isContractCallDescriptor
+import org.jetbrains.kotlin.descriptors.FunctionDescriptor
+import org.jetbrains.kotlin.descriptors.PackageFragmentDescriptor
+import org.jetbrains.kotlin.descriptors.PropertyAccessorDescriptor
+import org.jetbrains.kotlin.descriptors.isOverridable
+import org.jetbrains.kotlin.diagnostics.Errors
+import org.jetbrains.kotlin.psi.KtBlockExpression
+import org.jetbrains.kotlin.psi.KtDotQualifiedExpression
+import org.jetbrains.kotlin.psi.KtElement
+import org.jetbrains.kotlin.resolve.BindingContext
+import org.jetbrains.kotlin.resolve.calls.model.ResolvedCall
+import org.jetbrains.kotlin.resolve.scopes.LexicalScope
+import org.jetbrains.kotlin.resolve.scopes.LexicalScopeKind
+
+object ContractNotAllowedCallChecker : CallChecker {
+    override fun check(resolvedCall: ResolvedCall<*>, reportOn: PsiElement, context: CallCheckerContext) {
+        if (reportOn !is KtElement) return
+        val descriptor = resolvedCall.resultingDescriptor as? FunctionDescriptor ?: return
+        if (!descriptor.isContractCallDescriptor()) return
+
+        val callElement = resolvedCall.call.callElement
+        var hasErrors = false
+
+        fun contractNotAllowed(message: String) {
+            hasErrors = true
+            context.trace.report(Errors.CONTRACT_NOT_ALLOWED.on(reportOn, message))
+        }
+
+        val scope = context.scope
+        val functionDescriptor = scope.ownerDescriptor as? FunctionDescriptor
+
+        if (functionDescriptor == null || functionDescriptor is PropertyAccessorDescriptor)
+            contractNotAllowed("Contracts are allowed only for functions")
+
+        var inFunctionBodyBlock = true
+        if (scope.ownerDescriptor.containingDeclaration !is PackageFragmentDescriptor
+            || scope.kind != LexicalScopeKind.CODE_BLOCK
+            || (scope.parent as? LexicalScope)?.kind != LexicalScopeKind.FUNCTION_INNER_SCOPE
+        ) {
+            if (scope.kind == LexicalScopeKind.FUNCTION_INNER_SCOPE) {
+                contractNotAllowed("Contracts are allowed only in function body block")
+                inFunctionBodyBlock = false
+            } else {
+                contractNotAllowed("Contracts are allowed only for top-level functions")
+            }
+        }
+
+        if (functionDescriptor?.isOperator == true) contractNotAllowed("Contracts are not allowed for operator functions")
+
+        if (functionDescriptor?.isOverridable == true) contractNotAllowed("Contracts are not allowed for open functions")
+
+        var parent = callElement.parent
+        var element = callElement
+        if (parent is KtDotQualifiedExpression) {
+            element = parent
+            parent = parent.parent
+        }
+        if ((parent !is KtBlockExpression || parent.children.first { it is KtElement } != element) && inFunctionBodyBlock) {
+            contractNotAllowed("Contract should be the first statement")
+        }
+
+        if (hasErrors) {
+            context.trace.record(BindingContext.CONTRACT_NOT_ALLOWED, callElement, true)
+        }
+    }
+}

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

@@ -295,9 +295,8 @@ public class ExpressionTypingServices {
             }
             blockLevelVisitor = new ExpressionTypingVisitorDispatcher.ForBlock(expressionTypingComponents, annotationChecker, scope);
 
-            expressionTypingComponents.contractParsingServices.checkContractAndRecordIfPresent(statementExpression, context.trace, scope, isFirstStatement);
-
             if (isFirstStatement) {
+                expressionTypingComponents.contractParsingServices.checkContractAndRecordIfPresent(statementExpression, context.trace, scope);
                 isFirstStatement = false;
             }
         }