From 72b0d3132939dbd12c8d0647c75d76eafd738219 Mon Sep 17 00:00:00 2001
From: Dmitry Savvinov <dsavvinov@gmail.com>
Date: Tue, 3 Oct 2017 15:02:40 +0300
Subject: [PATCH] Effects: support contracts with conditional returns

Support functions which have explicitly expressed relation between
return-value and arguments. If we have observed that function returned
this value, get additional information from effect system.

==========
Effect System introduction: 10/18
---
 .../types/expressions/DataFlowAnalyzer.java      | 16 +++++++++++++---
 .../expressions/PatternMatchingTypingVisitor.kt  |  5 ++++-
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/compiler/frontend/src/org/jetbrains/kotlin/types/expressions/DataFlowAnalyzer.java b/compiler/frontend/src/org/jetbrains/kotlin/types/expressions/DataFlowAnalyzer.java
index ea24297d024..0beb1962d63 100644
--- a/compiler/frontend/src/org/jetbrains/kotlin/types/expressions/DataFlowAnalyzer.java
+++ b/compiler/frontend/src/org/jetbrains/kotlin/types/expressions/DataFlowAnalyzer.java
@@ -26,6 +26,7 @@ import org.jetbrains.kotlin.config.LanguageFeature;
 import org.jetbrains.kotlin.config.LanguageVersionSettings;
 import org.jetbrains.kotlin.descriptors.*;
 import org.jetbrains.kotlin.diagnostics.DiagnosticUtilsKt;
+import org.jetbrains.kotlin.contracts.EffectSystem;
 import org.jetbrains.kotlin.incremental.KotlinLookupLocation;
 import org.jetbrains.kotlin.lexer.KtTokens;
 import org.jetbrains.kotlin.psi.*;
@@ -59,6 +60,7 @@ public class DataFlowAnalyzer {
     private final SmartCastManager smartCastManager;
     private final ExpressionTypingFacade facade;
     private final LanguageVersionSettings languageVersionSettings;
+    private final EffectSystem effectSystem;
 
     public DataFlowAnalyzer(
             @NotNull Iterable<AdditionalTypeChecker> additionalTypeCheckers,
@@ -66,7 +68,8 @@ public class DataFlowAnalyzer {
             @NotNull KotlinBuiltIns builtIns,
             @NotNull SmartCastManager smartCastManager,
             @NotNull ExpressionTypingFacade facade,
-            @NotNull LanguageVersionSettings languageVersionSettings
+            @NotNull LanguageVersionSettings languageVersionSettings,
+            @NotNull EffectSystem effectSystem
     ) {
         this.additionalTypeCheckers = additionalTypeCheckers;
         this.constantExpressionEvaluator = constantExpressionEvaluator;
@@ -74,6 +77,7 @@ public class DataFlowAnalyzer {
         this.smartCastManager = smartCastManager;
         this.facade = facade;
         this.languageVersionSettings = languageVersionSettings;
+        this.effectSystem = effectSystem;
     }
 
     // NB: use this method only for functions from 'Any'
@@ -234,10 +238,16 @@ public class DataFlowAnalyzer {
                 }
             }
         });
+
+        DataFlowInfo infoFromEffectSystem = effectSystem.extractDataFlowInfoFromCondition(
+                condition, conditionValue, context.trace, DescriptorUtils.getContainingModule(context.scope.getOwnerDescriptor())
+        );
+
         if (result.get() == null) {
-            return context.dataFlowInfo;
+            return context.dataFlowInfo.and(infoFromEffectSystem);
         }
-        return context.dataFlowInfo.and(result.get());
+
+        return context.dataFlowInfo.and(result.get()).and(infoFromEffectSystem);
     }
 
     @Nullable
diff --git a/compiler/frontend/src/org/jetbrains/kotlin/types/expressions/PatternMatchingTypingVisitor.kt b/compiler/frontend/src/org/jetbrains/kotlin/types/expressions/PatternMatchingTypingVisitor.kt
index 7dee6a2167f..7ea9da30366 100644
--- a/compiler/frontend/src/org/jetbrains/kotlin/types/expressions/PatternMatchingTypingVisitor.kt
+++ b/compiler/frontend/src/org/jetbrains/kotlin/types/expressions/PatternMatchingTypingVisitor.kt
@@ -368,8 +368,11 @@ class PatternMatchingTypingVisitor internal constructor(facade: ExpressionTyping
             override fun visitWhenConditionWithExpression(condition: KtWhenConditionWithExpression) {
                 val expression = condition.expression
                 if (expression != null) {
-                    newDataFlowInfo = checkTypeForExpressionCondition(
+                    val basicDataFlowInfo = checkTypeForExpressionCondition(
                             context, expression, subjectType, subjectExpression == null, subjectDataFlowValue)
+                    val moduleDescriptor = DescriptorUtils.getContainingModule(context.scope.ownerDescriptor)
+                    val dataFlowInfoFromES = components.effectSystem.getDataFlowInfoWhenEquals(subjectExpression, expression, context.trace, moduleDescriptor)
+                    newDataFlowInfo = basicDataFlowInfo.and(dataFlowInfoFromES)
                 }
             }