[FIR] Unwrap when subject expression recursively in DFA variable storage

#KT-37327 Fixed

compiler/fir/resolve/src/org/jetbrains/kotlin/fir/resolve/dfa/VariableStorage.kt

@@ -42,7 +42,7 @@ class VariableStorage(val session: FirSession) {
     }
 
     private fun FirElement.unwrapElement(): FirElement = when (this) {
-        is FirWhenSubjectExpression -> whenSubject.whenExpression.let { it.subjectVariable ?: it.subject } ?: this
+        is FirWhenSubjectExpression -> whenSubject.whenExpression.let { it.subjectVariable ?: it.subject }?.unwrapElement() ?: this
         is FirExpressionWithSmartcast -> originalExpression.unwrapElement()
         else -> this
     }

compiler/fir/resolve/testData/resolve/smartcasts/kt37327.kt

@@ -0,0 +1,14 @@
+// ISSUE: KT-37327
+
+interface Q
+
+sealed class A : Q
+
+class B(val x: Int) : A()
+
+fun Q.foo() {
+    if (this !is A) return
+    when (this) {
+        is B -> x // unresolved
+    }
+}

compiler/fir/resolve/testData/resolve/smartcasts/kt37327.txt

@@ -0,0 +1,32 @@
+FILE: kt37327.kt
+    public abstract interface Q : R|kotlin/Any| {
+    }
+    public sealed class A : R|Q| {
+        private constructor(): R|A| {
+            super<R|kotlin/Any|>()
+        }
+
+    }
+    public final class B : R|A| {
+        public constructor(x: R|kotlin/Int|): R|B| {
+            super<R|A|>()
+        }
+
+        public final val x: R|kotlin/Int| = R|<local>/x|
+            public get(): R|kotlin/Int|
+
+    }
+    public final fun R|Q|.foo(): R|kotlin/Unit| {
+        when () {
+            (this@R|/foo| !is R|A|) ->  {
+                ^foo Unit
+            }
+        }
+
+        when (this@R|/foo|) {
+            ($subj$ is R|B|) ->  {
+                this@R|/foo|.R|/B.x|
+            }
+        }
+
+    }

compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirDiagnosticsTestGenerated.java

@@ -1574,6 +1574,11 @@ public class FirDiagnosticsTestGenerated extends AbstractFirDiagnosticsTest {
             runTest("compiler/fir/resolve/testData/resolve/smartcasts/equalsAndIdentity.kt");
         }
 
+        @TestMetadata("kt37327.kt")
+        public void testKt37327() throws Exception {
+            runTest("compiler/fir/resolve/testData/resolve/smartcasts/kt37327.kt");
+        }
+
         @TestMetadata("multipleCasts.kt")
         public void testMultipleCasts() throws Exception {
             runTest("compiler/fir/resolve/testData/resolve/smartcasts/multipleCasts.kt");

compiler/fir/resolve/tests/org/jetbrains/kotlin/fir/FirDiagnosticsWithLightTreeTestGenerated.java

@@ -1574,6 +1574,11 @@ public class FirDiagnosticsWithLightTreeTestGenerated extends AbstractFirDiagnos
             runTest("compiler/fir/resolve/testData/resolve/smartcasts/equalsAndIdentity.kt");
         }
 
+        @TestMetadata("kt37327.kt")
+        public void testKt37327() throws Exception {
+            runTest("compiler/fir/resolve/testData/resolve/smartcasts/kt37327.kt");
+        }
+
         @TestMetadata("multipleCasts.kt")
         public void testMultipleCasts() throws Exception {
             runTest("compiler/fir/resolve/testData/resolve/smartcasts/multipleCasts.kt");