Force-resolve body in contract computation under lock

Even though acquire/release pattern guarantees memory visibility across
threads, it doesn't prevents concurrent access to critical section (i.e.
to force-resolve of the corresponding body).

This can lead to multiple resolution passes over one and the same PSI in
IDE, which, in turn, leads to 'rewrite at slice'-exceptions. See
KT-30030 for case description and details.

^KT-30030 Fixed

compiler/frontend/src/org/jetbrains/kotlin/contracts/ContractDeserializerImpl.kt

@@ -41,14 +41,14 @@ class ContractDeserializerImpl(
         ownerFunction: FunctionDescriptor,
         typeTable: TypeTable,
         typeDeserializer: TypeDeserializer
-    ): Pair<CallableDescriptor.UserDataKey<*>, LazyContractProvider>? {
+    ): Pair<CallableDescriptor.UserDataKey<*>, AbstractContractProvider>? {
         if (!proto.hasContract()) return null
 
         if (!configuration.readDeserializedContracts) return null
 
         val worker = ContractDeserializationWorker(typeTable, typeDeserializer, ownerFunction, storageManager)
         val contract = worker.deserializeContract(proto.contract) ?: return null
-        return ContractProviderKey to LazyContractProvider.createInitialized(contract)
+        return ContractProviderKey to ContractProviderImpl(contract)
     }
 
     private class ContractDeserializationWorker(

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

@@ -21,6 +21,7 @@ import org.jetbrains.kotlin.config.LanguageFeature
 import org.jetbrains.kotlin.config.LanguageVersionSettings
 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.DeclarationDescriptor
 import org.jetbrains.kotlin.descriptors.FunctionDescriptor
 import org.jetbrains.kotlin.descriptors.PackageFragmentDescriptor
@@ -50,7 +51,8 @@ class ContractParsingServices(val languageVersionSettings: LanguageVersionSettin
         if (!expression.isContractDescriptionCallPsiCheck()) return
 
         val callContext = ContractCallContext(expression, isFirstStatement, scope, trace)
-        val contractProviderIfAny = (scope.ownerDescriptor as? FunctionDescriptor)?.getUserData(ContractProviderKey)
+        val contractProviderIfAny =
+            (scope.ownerDescriptor as? FunctionDescriptor)?.getUserData(ContractProviderKey) as? LazyContractProvider?
         var resultingContractDescription: ContractDescription? = null
 
         try {

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

@@ -263,7 +263,7 @@ class FunctionDescriptorResolver(
 
         if (!isContractsEnabled || !function.mayHaveContract()) return null
 
-        return LazyContractProvider {
+        return LazyContractProvider(storageManager) {
             AstLoadingFilter.forceAllowTreeLoading(function.containingFile) {
                 expressionTypingServices.getBodyExpressionType(trace, scope, dataFlowInfo, function, functionDescriptor)
             }

compiler/resolution/src/org/jetbrains/kotlin/contracts/description/LazyContractProvider.kt

@@ -18,21 +18,29 @@ package org.jetbrains.kotlin.contracts.description
 
 import org.jetbrains.kotlin.descriptors.CallableDescriptor
 import org.jetbrains.kotlin.serialization.deserialization.ContractProvider
+import org.jetbrains.kotlin.storage.StorageManager
+
+abstract class AbstractContractProvider : ContractProvider {
+    abstract fun getContractDescription(): ContractDescription?
+}
 
 /**
  * Essentially, this is a composition of two fields: value of type 'ContractDescription' and
  * 'computation', which guarantees to initialize this field.
+ *
+ * Such contract providers are present only for source-based declarations, where we need additional
+ * resolve (force-resolve of the body) to get ContractDescription
  */
-class LazyContractProvider(private val computation: () -> Any?) : ContractProvider {
+class LazyContractProvider(private val storageManager: StorageManager, private val computation: () -> Any?) : AbstractContractProvider() {
     @Volatile
     private var isComputed: Boolean = false
 
     private var contractDescription: ContractDescription? = null
 
 
-    fun getContractDescription(): ContractDescription? {
+    override fun getContractDescription(): ContractDescription? {
         if (!isComputed) {
-            computation.invoke() // should initialize contractDescription
+            storageManager.compute(computation)
             assert(isComputed) { "Computation of contract hasn't initialized contract properly" }
         }
 
@@ -43,12 +51,15 @@ class LazyContractProvider(private val computation: () -> Any?) : ContractProvid
         this.contractDescription = contractDescription
         isComputed = true // publish
     }
+}
 
-    companion object {
-        fun createInitialized(contract: ContractDescription?): LazyContractProvider =
-            LazyContractProvider({}).apply { setContractDescription(contract) }
-    }
+/**
+ * Such contract providers are used where we can be sure about contract presence and don't need
+ * additional resolve (e.g., for deserialized declarations)
+ */
+class ContractProviderImpl(private val contractDescription: ContractDescription) : AbstractContractProvider() {
+    override fun getContractDescription(): ContractDescription = contractDescription
 }
 
 // For storing into UserDataMap of FunctionDescriptor
-object ContractProviderKey : CallableDescriptor.UserDataKey<LazyContractProvider?>
+object ContractProviderKey : CallableDescriptor.UserDataKey<AbstractContractProvider?>

compiler/tests-common/tests/org/jetbrains/kotlin/test/util/RecursiveDescriptorComparator.java

@@ -201,7 +201,7 @@ public class RecursiveDescriptorComparator {
     }
 
     private static void printEffectsIfAny(FunctionDescriptor functionDescriptor, Printer printer) {
-        LazyContractProvider contractProvider = functionDescriptor.getUserData(ContractProviderKey.INSTANCE);
+        AbstractContractProvider contractProvider = functionDescriptor.getUserData(ContractProviderKey.INSTANCE);
         if (contractProvider == null) return;
 
         ContractDescription contractDescription = contractProvider.getContractDescription();