Expand Effect System with receiver values (#KT-28672)

Add some classes to hierarchy of `ESValue`:
- `ESReceiver` -- representation of receiver
  of lambda function
- `ESDataFlow` and `ESDataFlowReceiver` -- little refactoring of entities
  that holds information for DFA (description in comments in code).

Also add kdoc to classes of `ESValue` hierarchy

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

@@ -16,31 +16,67 @@
 
 package org.jetbrains.kotlin.contracts
 
-import org.jetbrains.kotlin.contracts.model.ESValue
-import org.jetbrains.kotlin.contracts.model.structure.ESVariable
 import org.jetbrains.kotlin.contracts.model.ESExpressionVisitor
+import org.jetbrains.kotlin.contracts.model.structure.AbstractESValue
+import org.jetbrains.kotlin.contracts.model.structure.ESReceiverValue
+import org.jetbrains.kotlin.contracts.model.structure.ESVariable
 import org.jetbrains.kotlin.descriptors.ValueDescriptor
 import org.jetbrains.kotlin.psi.KtLambdaExpression
 import org.jetbrains.kotlin.resolve.calls.smartcasts.DataFlowValue
+import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue
 
-class ESDataFlowValue(descriptor: ValueDescriptor, val dataFlowValue: DataFlowValue) : ESVariable(descriptor) {
-    override fun equals(other: Any?): Boolean {
+
+/**
+ * [ESDataFlowValue] is an interface, representing some entity that holds [DataFlowValue] for DFA.
+ *
+ * Actually that interfaces must be sealed.
+ */
+interface ESDataFlowValue {
+    val dataFlowValue: DataFlowValue
+
+    fun dataFlowEquals(other: Any?): Boolean {
         if (this === other) return true
-        if (javaClass != other?.javaClass) return false
+        if (other !is ESDataFlowValue) return false
 
-        other as ESDataFlowValue
-
-        if (dataFlowValue != other.dataFlowValue) return false
-
-        return true
+        return dataFlowValue == other.dataFlowValue
     }
+}
+
+
+/**
+ * [ESVariableWithDataFlowValue] is [ESVariable] with data flow information.
+ */
+class ESVariableWithDataFlowValue(
+    descriptor: ValueDescriptor,
+    override val dataFlowValue: DataFlowValue
+) : ESVariable(descriptor), ESDataFlowValue {
+    override fun equals(other: Any?): Boolean = dataFlowEquals(other)
 
     override fun hashCode(): Int {
         return dataFlowValue.hashCode()
     }
 }
 
-class ESLambda(val lambda: KtLambdaExpression) : ESValue(null) {
+
+/**
+ * [ESReceiverWithDataFlowValue] is [ESReceiverValue] with data flow information.
+ */
+class ESReceiverWithDataFlowValue(
+    receiverValue: ReceiverValue,
+    override val dataFlowValue: DataFlowValue
+) : ESReceiverValue(receiverValue), ESDataFlowValue {
+    override fun equals(other: Any?): Boolean = dataFlowEquals(other)
+
+    override fun hashCode(): Int {
+        return dataFlowValue.hashCode()
+    }
+}
+
+
+/**
+ * [ESLambda] represents lambda functions in Effect System
+ */
+class ESLambda(val lambda: KtLambdaExpression) : AbstractESValue(null) {
     override fun <T> accept(visitor: ESExpressionVisitor<T>): T {
         throw IllegalStateException("Lambdas shouldn't be visited by ESExpressionVisitor")
     }

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

@@ -43,6 +43,7 @@ import org.jetbrains.kotlin.resolve.calls.tasks.ExplicitReceiverKind
 import org.jetbrains.kotlin.resolve.constants.CompileTimeConstant
 import org.jetbrains.kotlin.resolve.constants.UnsignedErrorValueTypeConstant
 import org.jetbrains.kotlin.resolve.scopes.receivers.ExpressionReceiver
+import org.jetbrains.kotlin.resolve.scopes.receivers.ExtensionReceiver
 import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.utils.addIfNotNull
@@ -76,7 +77,7 @@ class EffectsExtractingVisitor(
                 builtIns.booleanType,
                 EqualsFunctor(constants, false).invokeWithArguments(arguments)
             )
-            descriptor is ValueDescriptor -> ESDataFlowValue(
+            descriptor is ValueDescriptor -> ESVariableWithDataFlowValue(
                 descriptor,
                 (element as KtExpression).createDataFlowValue() ?: return UNKNOWN_COMPUTATION
             )
@@ -161,9 +162,18 @@ class EffectsExtractingVisitor(
 
     private fun ReceiverValue.toComputation(): Computation = when (this) {
         is ExpressionReceiver -> extractOrGetCached(expression)
+        is ExtensionReceiver -> ESReceiverWithDataFlowValue(this, createDataFlowValue())
         else -> UNKNOWN_COMPUTATION
     }
 
+    private fun ExtensionReceiver.createDataFlowValue(): DataFlowValue {
+        return dataFlowValueFactory.createDataFlowValue(
+            receiverValue = this,
+            bindingContext = trace.bindingContext,
+            containingDeclarationOrModule = this.declarationDescriptor
+        )
+    }
+
     private fun KtExpression.createDataFlowValue(): DataFlowValue? {
         return dataFlowValueFactory.createDataFlowValue(
             expression = this,

compiler/resolution/src/org/jetbrains/kotlin/contracts/model/ESExpressionVisitor.kt

@@ -28,4 +28,6 @@ interface ESExpressionVisitor<out T> {
 
     fun visitVariable(esVariable: ESVariable): T
     fun visitConstant(esConstant: ESConstant): T
+
+    fun visitReceiver(esReceiver: ESReceiver): T
 }

compiler/resolution/src/org/jetbrains/kotlin/contracts/model/ESExpressions.kt

@@ -16,16 +16,21 @@
 
 package org.jetbrains.kotlin.contracts.model
 
-import org.jetbrains.kotlin.types.KotlinType
 
+/**
+ * There is a subset of Kotlin language in Effect system (expressions
+ *   in right hand side of conditional effect) and [ESExpression] with subtypes
+ *   precisely enumerate what can be found here.
+ */
 interface ESExpression {
     fun <T> accept(visitor: ESExpressionVisitor<T>): T
 }
 
 interface ESOperator : ESExpression {
+    /**
+     * [Functor] that contains logic of concrete operator
+     */
     val functor: Functor
 }
 
-abstract class ESValue(override val type: KotlinType?) : Computation, ESExpression {
-    override val effects: List<ESEffect> = listOf()
-}
+interface ESValue : Computation, ESExpression

compiler/resolution/src/org/jetbrains/kotlin/contracts/model/structure/Values.kt

@@ -19,14 +19,45 @@ package org.jetbrains.kotlin.contracts.model.structure
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.contracts.description.expressions.BooleanConstantReference
 import org.jetbrains.kotlin.contracts.description.expressions.ConstantReference
+import org.jetbrains.kotlin.contracts.model.ESEffect
 import org.jetbrains.kotlin.contracts.model.ESExpression
 import org.jetbrains.kotlin.contracts.model.ESExpressionVisitor
 import org.jetbrains.kotlin.contracts.model.ESValue
 import org.jetbrains.kotlin.descriptors.ValueDescriptor
+import org.jetbrains.kotlin.resolve.scopes.receivers.ReceiverValue
 import org.jetbrains.kotlin.types.KotlinType
 import java.util.*
 
-open class ESVariable(val descriptor: ValueDescriptor) : ESValue(descriptor.type) {
+
+interface ESReceiver : ESValue {
+    val receiverValue: ReceiverValue
+
+    override fun <T> accept(visitor: ESExpressionVisitor<T>): T = visitor.visitReceiver(this)
+}
+
+
+abstract class AbstractESValue(override val type: KotlinType?) : ESValue {
+    override val effects: List<ESEffect> = listOf()
+}
+
+
+open class ESReceiverValue(override val receiverValue: ReceiverValue) : AbstractESValue(null), ESReceiver
+
+
+/**
+ * [ESVariable] is class with multiple applications.
+ *
+ * 1. [ESVariable] represents some variable on declaration-site of contract (reference to parameter
+ *   of function). @see [org.jetbrains.kotlin.contracts.interpretation.ContractInterpretationDispatcher.interpretVariable].
+ * 2. [ESVariable] is wrapper around argument passed to function in process of substitution.
+ *   @see [org.jetbrains.kotlin.contracts.EffectsExtractingVisitor.visitKtElement].
+ * 3. [ESVariable] is a key in [Substitutor], that maps values from function signature to
+ *   real values from call-site. That keys are equal to variables from point 1.
+ *   @see [org.jetbrains.kotlin.contracts.model.functors.SubstitutingFunctor.doInvocation].
+ *
+ * [ESVariable] at points 2 and 3 must has consistent equality according to using them as keys
+ */
+open class ESVariable(val descriptor: ValueDescriptor) : AbstractESValue(descriptor.type) {
     override fun <T> accept(visitor: ESExpressionVisitor<T>): T = visitor.visitVariable(this)
 
     override fun equals(other: Any?): Boolean {
@@ -45,7 +76,13 @@ open class ESVariable(val descriptor: ValueDescriptor) : ESValue(descriptor.type
     override fun toString(): String = descriptor.toString()
 }
 
-class ESConstant internal constructor(val constantReference: ConstantReference, override val type: KotlinType) : ESValue(type) {
+
+/**
+ * [ESConstant] represent some constant is Effect System
+ *
+ * There is only few constants are supported (@see [ESConstant.Companion])
+ */
+class ESConstant internal constructor(val constantReference: ConstantReference, override val type: KotlinType) : AbstractESValue(type) {
     override fun <T> accept(visitor: ESExpressionVisitor<T>): T = visitor.visitConstant(this)
 
     override fun equals(other: Any?): Boolean = other is ESConstant && constantReference == other.constantReference

compiler/resolution/src/org/jetbrains/kotlin/contracts/model/visitors/InfoCollector.kt

@@ -83,6 +83,8 @@ class InfoCollector(
 
     override fun visitConstant(esConstant: ESConstant): MutableContextInfo = MutableContextInfo.EMPTY
 
+    override fun visitReceiver(esReceiver: ESReceiver): MutableContextInfo = MutableContextInfo.EMPTY
+
     private fun <R> inverted(block: () -> R): R {
         isInverted = isInverted.not()
         val result = block()

compiler/resolution/src/org/jetbrains/kotlin/contracts/model/visitors/Reducer.kt

@@ -30,7 +30,6 @@ class Reducer(private val constants: ESConstants) : ESExpressionVisitor<ESExpres
 
     private fun reduceEffect(effect: ESEffect): ESEffect? {
         when (effect) {
-            is SimpleEffect -> return effect
             is ConditionalEffect -> {
                 // Reduce condition
                 val reducedCondition = effect.condition.accept(this) ?: return null
@@ -44,6 +43,7 @@ class Reducer(private val constants: ESConstants) : ESExpressionVisitor<ESExpres
                 // Leave everything else as is
                 return effect
             }
+            else -> return effect
         }
     }
 
@@ -108,4 +108,6 @@ class Reducer(private val constants: ESConstants) : ESExpressionVisitor<ESExpres
     override fun visitVariable(esVariable: ESVariable): ESVariable = esVariable
 
     override fun visitConstant(esConstant: ESConstant): ESConstant = esConstant
+
+    override fun visitReceiver(esReceiver: ESReceiver): ESReceiver = esReceiver
 }

compiler/resolution/src/org/jetbrains/kotlin/contracts/model/visitors/Substitutor.kt

@@ -62,4 +62,6 @@ class Substitutor(
     override fun visitVariable(esVariable: ESVariable): Computation? = substitutions[esVariable] ?: esVariable
 
     override fun visitConstant(esConstant: ESConstant): Computation? = esConstant
+
+    override fun visitReceiver(esReceiver: ESReceiver): ESReceiver = esReceiver
 }

compiler/testData/diagnostics/testsWithStdLib/contracts/smartcasts/extensionReceiver.kt

@@ -20,7 +20,7 @@ fun smartcastOnReceiver(s: String?) {
             <!UNSAFE_CALL!>length<!>
         }
         else {
-            <!UNSAFE_CALL!>length<!>
+            <!DEBUG_INFO_IMPLICIT_RECEIVER_SMARTCAST!>length<!>
         }
     }
 }