From 724b7bf363e33375fe7c2f1cdef229a81434dbb0 Mon Sep 17 00:00:00 2001
From: Dmitry Savvinov <dmitry.savvinov@jetbrains.com>
Date: Thu, 19 Jul 2018 19:06:21 +0300
Subject: [PATCH] Minor: rename and refactor contract-related PSI checks in
 ktPsiUtil

---
 .../parsing/ContractParsingServices.kt        | 16 +++++++-------
 .../resolve/FunctionDescriptorResolver.kt     |  4 ++--
 .../jetbrains/kotlin/psi/psiUtil/ktPsiUtil.kt | 21 +++++++++++--------
 3 files changed, 22 insertions(+), 19 deletions(-)

diff --git a/compiler/frontend/src/org/jetbrains/kotlin/contracts/parsing/ContractParsingServices.kt b/compiler/frontend/src/org/jetbrains/kotlin/contracts/parsing/ContractParsingServices.kt
index a175a858855..ef96c35fe1d 100644
--- a/compiler/frontend/src/org/jetbrains/kotlin/contracts/parsing/ContractParsingServices.kt
+++ b/compiler/frontend/src/org/jetbrains/kotlin/contracts/parsing/ContractParsingServices.kt
@@ -24,7 +24,7 @@ import org.jetbrains.kotlin.contracts.description.ContractProviderKey
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.diagnostics.Errors
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.isContractDescriptionCallFastCheck
+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
@@ -34,7 +34,7 @@ import org.jetbrains.kotlin.resolve.scopes.LexicalScopeKind
 class ContractParsingServices(val languageVersionSettings: LanguageVersionSettings) {
     fun checkContractAndRecordIfPresent(expression: KtExpression, trace: BindingTrace, scope: LexicalScope, isFirstStatement: Boolean) {
         val ownerDescriptor = scope.ownerDescriptor
-        if (!isContractDescriptionCallFastCheck(expression) || ownerDescriptor !is FunctionDescriptor) return
+        if (!expression.isContractDescriptionCallPsiCheck() || ownerDescriptor !is FunctionDescriptor) return
         val contractProvider = ownerDescriptor.getUserData(ContractProviderKey) ?: return
 
         val isFeatureTurnedOn = languageVersionSettings.supportsFeature(LanguageFeature.AllowContractsForCustomFunctions) ||
@@ -42,7 +42,7 @@ class ContractParsingServices(val languageVersionSettings: LanguageVersionSettin
                 languageVersionSettings.getFlag(AnalysisFlag.Flags.allowKotlinPackage)
 
         val contractDescriptor = when {
-            !isContractDescriptionCallPreciseCheck(expression, trace.bindingContext) -> null
+            !expression.isContractDescriptionCallPreciseCheck(trace.bindingContext) -> null
 
             !isFeatureTurnedOn -> {
                 trace.report(
@@ -65,15 +65,15 @@ class ContractParsingServices(val languageVersionSettings: LanguageVersionSettin
         contractProvider.setContractDescription(contractDescriptor)
     }
 
+    internal fun isContractDescriptionCall(expression: KtExpression, context: BindingContext): Boolean =
+        expression.isContractDescriptionCallPsiCheck() && expression.isContractDescriptionCallPreciseCheck(context)
+
     private fun parseContract(expression: KtExpression?, trace: BindingTrace, ownerDescriptor: FunctionDescriptor): ContractDescription? =
         PsiContractParserDispatcher(trace, this).parseContract(expression, ownerDescriptor)
 
-    internal fun isContractDescriptionCall(expression: KtExpression, context: BindingContext): Boolean =
-        isContractDescriptionCallFastCheck(expression) && isContractDescriptionCallPreciseCheck(expression, context)
-
     private fun isContractAllowedHere(scope: LexicalScope): Boolean =
         scope.kind == LexicalScopeKind.CODE_BLOCK && (scope.parent as? LexicalScope)?.kind == LexicalScopeKind.FUNCTION_INNER_SCOPE
 
-    private fun isContractDescriptionCallPreciseCheck(expression: KtExpression, context: BindingContext): Boolean =
-        expression.getResolvedCall(context)?.resultingDescriptor?.isContractCallDescriptor() ?: false
+    private fun KtExpression.isContractDescriptionCallPreciseCheck(context: BindingContext): Boolean =
+        getResolvedCall(context)?.resultingDescriptor?.isContractCallDescriptor() ?: false
 }
\ No newline at end of file
diff --git a/compiler/frontend/src/org/jetbrains/kotlin/resolve/FunctionDescriptorResolver.kt b/compiler/frontend/src/org/jetbrains/kotlin/resolve/FunctionDescriptorResolver.kt
index 6898b7f9db7..07914dd6b71 100644
--- a/compiler/frontend/src/org/jetbrains/kotlin/resolve/FunctionDescriptorResolver.kt
+++ b/compiler/frontend/src/org/jetbrains/kotlin/resolve/FunctionDescriptorResolver.kt
@@ -39,7 +39,7 @@ import org.jetbrains.kotlin.diagnostics.PsiDiagnosticUtils
 import org.jetbrains.kotlin.lexer.KtTokens
 import org.jetbrains.kotlin.name.Name
 import org.jetbrains.kotlin.psi.*
-import org.jetbrains.kotlin.psi.psiUtil.fastCheckIfContractPresent
+import org.jetbrains.kotlin.psi.psiUtil.isContractPresentPsiCheck
 import org.jetbrains.kotlin.psi.psiUtil.hasActualModifier
 import org.jetbrains.kotlin.psi.psiUtil.hasExpectModifier
 import org.jetbrains.kotlin.resolve.DescriptorResolver.getDefaultModality
@@ -242,7 +242,7 @@ class FunctionDescriptorResolver(
                 // We need to enable contracts if we're compiling "kotlin"-package to be able to ship contracts in stdlib in 1.2
                 languageVersionSettings.getFlag(AnalysisFlag.allowKotlinPackage)
 
-        if (!isContractsEnabled || !fastCheckIfContractPresent(function)) return emptyContract
+        if (!isContractsEnabled || !function.isContractPresentPsiCheck()) return emptyContract
 
         return provideByDeferredForceResolve
     }
diff --git a/compiler/psi/src/org/jetbrains/kotlin/psi/psiUtil/ktPsiUtil.kt b/compiler/psi/src/org/jetbrains/kotlin/psi/psiUtil/ktPsiUtil.kt
index 5ae3a9eb7d9..9e7f1e45eed 100644
--- a/compiler/psi/src/org/jetbrains/kotlin/psi/psiUtil/ktPsiUtil.kt
+++ b/compiler/psi/src/org/jetbrains/kotlin/psi/psiUtil/ktPsiUtil.kt
@@ -293,17 +293,20 @@ inline fun <reified T : KtElement, R> flatMapDescendantsOfTypeVisitor(
 
 // ----------- Contracts -------------------------------------------------------------------------------------------------------------------
 
-fun fastCheckIfContractPresent(element: KtElement): Boolean {
-    if (!isContractAllowedHere(element)) return false
-    val firstExpression = ((element as? KtFunction)?.bodyExpression as? KtBlockExpression)?.statements?.firstOrNull() ?: return false
-    return isContractDescriptionCallFastCheck(firstExpression)
+fun KtElement.isContractPresentPsiCheck(): Boolean {
+    val contractAllowedHere = this is KtNamedFunction &&
+            isTopLevel &&
+            hasBlockBody() &&
+            !hasModifier(KtTokens.OPERATOR_KEYWORD)
+    if (!contractAllowedHere) return false
+
+    val firstExpression = ((this as? KtFunction)?.bodyExpression as? KtBlockExpression)?.statements?.firstOrNull() ?: return false
+
+    return firstExpression.isContractDescriptionCallPsiCheck()
 }
 
-private fun isContractAllowedHere(element: KtElement): Boolean =
-    element is KtNamedFunction && element.isTopLevel && element.hasBlockBody() && !element.hasModifier(KtTokens.OPERATOR_KEYWORD)
-
-fun isContractDescriptionCallFastCheck(expression: KtExpression): Boolean =
-    expression is KtCallExpression && expression.calleeExpression?.text == "contract"
+fun KtExpression.isContractDescriptionCallPsiCheck(): Boolean =
+    this is KtCallExpression && calleeExpression?.text == "contract"
 
 
 // ----------- Other -----------------------------------------------------------------------------------------------------------------------