From cf61a6c30f36a4f6645b72dc7462efc31fb5baae Mon Sep 17 00:00:00 2001
From: Dmitriy Novozhilov <dmitriy.novozhilov@jetbrains.com>
Date: Mon, 11 Nov 2019 11:45:25 +0300
Subject: [PATCH] [FIR] Add support of multiple condition branches to
 exhaustiveness checker

---
 .../FirWhenExhaustivenessTransformer.kt           | 12 ++++++++++++
 .../testData/resolve/exhaustiveness_enum.kt       |  7 +++++++
 .../testData/resolve/exhaustiveness_enum.txt      | 11 +++++++++++
 .../resolve/exhaustiveness_sealedClass.kt         | 11 +++++++++--
 .../resolve/exhaustiveness_sealedClass.txt        | 15 +++++++++++++--
 5 files changed, 52 insertions(+), 4 deletions(-)

diff --git a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirWhenExhaustivenessTransformer.kt b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirWhenExhaustivenessTransformer.kt
index 6ea60f68967..77e82b5fd48 100644
--- a/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirWhenExhaustivenessTransformer.kt
+++ b/compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/transformers/FirWhenExhaustivenessTransformer.kt
@@ -101,6 +101,12 @@ class FirWhenExhaustivenessTransformer(private val bodyResolveComponents: BodyRe
                 }
             }
         }
+
+        override fun visitBinaryLogicExpression(binaryLogicExpression: FirBinaryLogicExpression, data: EnumExhaustivenessData) {
+            if (binaryLogicExpression.kind == LogicOperationKind.OR) {
+                binaryLogicExpression.acceptChildren(this, data)
+            }
+        }
     }
 
     // ------------------------ Sealed class exhaustiveness ------------------------
@@ -147,6 +153,12 @@ class FirWhenExhaustivenessTransformer(private val bodyResolveComponents: BodyRe
             val symbol = data.symbolProvider.getSymbolByLookupTag(lookupTag) as? FirClassSymbol ?: return
             data.visitedInheritors.replace(symbol.classId, true)
         }
+
+        override fun visitBinaryLogicExpression(binaryLogicExpression: FirBinaryLogicExpression, data: SealedExhaustivenessData) {
+            if (binaryLogicExpression.kind == LogicOperationKind.OR) {
+                binaryLogicExpression.acceptChildren(this, data)
+            }
+        }
     }
 
     // ------------------------ Boolean exhaustiveness ------------------------
diff --git a/compiler/fir/resolve/testData/resolve/exhaustiveness_enum.kt b/compiler/fir/resolve/testData/resolve/exhaustiveness_enum.kt
index 953cf3446a0..da4c4b24aa6 100644
--- a/compiler/fir/resolve/testData/resolve/exhaustiveness_enum.kt
+++ b/compiler/fir/resolve/testData/resolve/exhaustiveness_enum.kt
@@ -46,4 +46,11 @@ fun test_2(e: Enum?) {
         Enum.C -> 3
         else -> 4
     }
+}
+
+fun test_3(e: Enum) {
+    val a = when (e) {
+        Enum.A, Enum.B -> 1
+        Enum.C -> 2
+    }
 }
\ No newline at end of file
diff --git a/compiler/fir/resolve/testData/resolve/exhaustiveness_enum.txt b/compiler/fir/resolve/testData/resolve/exhaustiveness_enum.txt
index bbba5d6968e..755ba71b30a 100644
--- a/compiler/fir/resolve/testData/resolve/exhaustiveness_enum.txt
+++ b/compiler/fir/resolve/testData/resolve/exhaustiveness_enum.txt
@@ -120,3 +120,14 @@ FILE: exhaustiveness_enum.kt
         }
 
     }
+    public final fun test_3(e: R|Enum|): R|kotlin/Unit| {
+        lval a: R|kotlin/Int| = when (R|<local>/e|) {
+            ==($subj$, Q|Enum.A|) || ==($subj$, Q|Enum.B|) ->  {
+                Int(1)
+            }
+            ==($subj$, Q|Enum.C|) ->  {
+                Int(2)
+            }
+        }
+
+    }
diff --git a/compiler/fir/resolve/testData/resolve/exhaustiveness_sealedClass.kt b/compiler/fir/resolve/testData/resolve/exhaustiveness_sealedClass.kt
index b903f63da44..8ab05066465 100644
--- a/compiler/fir/resolve/testData/resolve/exhaustiveness_sealedClass.kt
+++ b/compiler/fir/resolve/testData/resolve/exhaustiveness_sealedClass.kt
@@ -37,17 +37,24 @@ fun test_2(e: Base?) {
         is C -> 3
     }
 
-    val a = when (e) {
+    val b = when (e) {
         is Base.A -> 1
         is Base.A.B -> 2
         is C -> 3
         null -> 4
     }
 
-    val a = when (e) {
+    val c = when (e) {
         is Base.A -> 1
         is Base.A.B -> 2
         is C -> 3
         else -> 4
     }
+}
+
+fun test_3(e: Base) {
+    val a = when (e) {
+        is Base.A, is Base.A.B -> 1
+        is C -> 2
+    }
 }
\ No newline at end of file
diff --git a/compiler/fir/resolve/testData/resolve/exhaustiveness_sealedClass.txt b/compiler/fir/resolve/testData/resolve/exhaustiveness_sealedClass.txt
index dcc4a0f1f16..723c1e4b13e 100644
--- a/compiler/fir/resolve/testData/resolve/exhaustiveness_sealedClass.txt
+++ b/compiler/fir/resolve/testData/resolve/exhaustiveness_sealedClass.txt
@@ -82,7 +82,7 @@ FILE: exhaustiveness_sealedClass.kt
             }
         }
 
-        lval a: R|kotlin/Int| = when (R|<local>/e|) {
+        lval b: R|kotlin/Int| = when (R|<local>/e|) {
             ($subj$ is R|Base.A|) ->  {
                 Int(1)
             }
@@ -97,7 +97,7 @@ FILE: exhaustiveness_sealedClass.kt
             }
         }
 
-        lval a: R|kotlin/Int| = when (R|<local>/e|) {
+        lval c: R|kotlin/Int| = when (R|<local>/e|) {
             ($subj$ is R|Base.A|) ->  {
                 Int(1)
             }
@@ -113,3 +113,14 @@ FILE: exhaustiveness_sealedClass.kt
         }
 
     }
+    public final fun test_3(e: R|Base|): R|kotlin/Unit| {
+        lval a: R|kotlin/Int| = when (R|<local>/e|) {
+            ($subj$ is R|Base.A|) || ($subj$ is R|Base.A.B|) ->  {
+                Int(1)
+            }
+            ($subj$ is R|C|) ->  {
+                Int(2)
+            }
+        }
+
+    }