diff --git a/compiler/backend/src/org/jetbrains/kotlin/codegen/ExpressionCodegen.java b/compiler/backend/src/org/jetbrains/kotlin/codegen/ExpressionCodegen.java
index 2d1318429aa..baf5eae07d8 100644
--- a/compiler/backend/src/org/jetbrains/kotlin/codegen/ExpressionCodegen.java
+++ b/compiler/backend/src/org/jetbrains/kotlin/codegen/ExpressionCodegen.java
@@ -4044,12 +4044,9 @@ The "returned" value of try expression with no finally is either the last expres
     }
 
     private boolean isExhaustive(@NotNull KtWhenExpression whenExpression, boolean isStatement) {
-        if (isStatement) {
-            return Boolean.TRUE.equals(bindingContext.get(BindingContext.IMPLICIT_EXHAUSTIVE_WHEN, whenExpression));
-        }
-        else {
-            return Boolean.TRUE.equals(bindingContext.get(BindingContext.EXHAUSTIVE_WHEN, whenExpression));
-        }
+        boolean exhaustive = Boolean.TRUE.equals(bindingContext.get(BindingContext.EXHAUSTIVE_WHEN, whenExpression));
+        return isStatement ? exhaustive || Boolean.TRUE.equals(bindingContext.get(BindingContext.IMPLICIT_EXHAUSTIVE_WHEN, whenExpression))
+                           : exhaustive;
     }
 
     public void putUnitInstanceOntoStackForNonExhaustiveWhen(
diff --git a/compiler/testData/codegen/box/when/implicitExhaustiveAndReturn.kt b/compiler/testData/codegen/box/when/implicitExhaustiveAndReturn.kt
new file mode 100644
index 00000000000..d28e7064678
--- /dev/null
+++ b/compiler/testData/codegen/box/when/implicitExhaustiveAndReturn.kt
@@ -0,0 +1,9 @@
+fun test(i: Int): String {
+    when (i) {
+        0 -> return "0"
+        1 -> return "1"
+    }
+    return "OK"
+}
+
+fun box(): String = test(42)
\ No newline at end of file
diff --git a/compiler/testData/codegen/bytecodeText/when/exhaustiveWhenUnit.kt b/compiler/testData/codegen/bytecodeText/when/exhaustiveWhenUnit.kt
new file mode 100644
index 00000000000..17f33e1fa5f
--- /dev/null
+++ b/compiler/testData/codegen/bytecodeText/when/exhaustiveWhenUnit.kt
@@ -0,0 +1,18 @@
+enum class AccessMode { READ, WRITE, EXECUTE }
+
+fun whenExpr(access: AccessMode) {
+    return when (access) {
+        AccessMode.READ -> {}
+        AccessMode.WRITE -> {}
+        AccessMode.EXECUTE -> {}
+    }
+}
+
+fun box(): String {
+    whenExpr(AccessMode.EXECUTE)
+    return "OK"
+}
+
+// 1 TABLESWITCH
+// 0 LOOKUPSWITCH
+// 1 ATHROW
diff --git a/compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java b/compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java
index 1e4e302cf25..8996ae6dc0e 100644
--- a/compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java
+++ b/compiler/tests/org/jetbrains/kotlin/codegen/BlackBoxCodegenTestGenerated.java
@@ -14524,6 +14524,12 @@ public class BlackBoxCodegenTestGenerated extends AbstractBlackBoxCodegenTest {
             doTest(fileName);
         }
 
+        @TestMetadata("implicitExhaustiveAndReturn.kt")
+        public void testImplicitExhaustiveAndReturn() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/implicitExhaustiveAndReturn.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("integralWhenWithNoInlinedConstants.kt")
         public void testIntegralWhenWithNoInlinedConstants() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/box/when/integralWhenWithNoInlinedConstants.kt");
diff --git a/compiler/tests/org/jetbrains/kotlin/codegen/BytecodeTextTestGenerated.java b/compiler/tests/org/jetbrains/kotlin/codegen/BytecodeTextTestGenerated.java
index 27d35746b74..e239c72d40d 100644
--- a/compiler/tests/org/jetbrains/kotlin/codegen/BytecodeTextTestGenerated.java
+++ b/compiler/tests/org/jetbrains/kotlin/codegen/BytecodeTextTestGenerated.java
@@ -1294,6 +1294,12 @@ public class BytecodeTextTestGenerated extends AbstractBytecodeTextTest {
             doTest(fileName);
         }
 
+        @TestMetadata("exhaustiveWhenUnit.kt")
+        public void testExhaustiveWhenUnit() throws Exception {
+            String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeText/when/exhaustiveWhenUnit.kt");
+            doTest(fileName);
+        }
+
         @TestMetadata("integralWhenWithNoInlinedConstants.kt")
         public void testIntegralWhenWithNoInlinedConstants() throws Exception {
             String fileName = KotlinTestUtils.navigationMetadata("compiler/testData/codegen/bytecodeText/when/integralWhenWithNoInlinedConstants.kt");