Narrow references within optional chains in for-in loops (microsoft#52059)

Author: Andarist

src/compiler/checker.ts

@@ -26273,7 +26273,11 @@ export function createTypeChecker(host: TypeCheckerHost): TypeChecker {
                 return declaredType;
             }
             // for (const _ in ref) acts as a nonnull on ref
-            if (isVariableDeclaration(node) && node.parent.parent.kind === SyntaxKind.ForInStatement && isMatchingReference(reference, node.parent.parent.expression)) {
+            if (
+                isVariableDeclaration(node) &&
+                node.parent.parent.kind === SyntaxKind.ForInStatement &&
+                (isMatchingReference(reference, node.parent.parent.expression) || optionalChainContainsReference(node.parent.parent.expression, reference))
+            ) {
                 return getNonNullableTypeIfNeeded(finalizeEvolvingArrayType(getTypeFromFlowType(getTypeAtFlowNode(flow.antecedent))));
             }
             // Assignment doesn't affect reference

tests/baselines/reference/controlFlowOptionalChain.errors.txt

@@ -773,4 +773,21 @@ tests/cases/conformance/controlFlow/controlFlowOptionalChain.ts(567,21): error T
             console.log("I should ALSO be reachable");
         }
     }
-    
+    
+    
+    // Repro from #51941
+    
+    type Test5 = {
+      main?: {
+        childs: Record<string, Test5>;
+      };
+    };
+    
+    function f50(obj: Test5) {
+       for (const key in obj.main?.childs) {
+          if (obj.main.childs[key] === obj) {
+            return obj;
+          }
+       }
+       return null;
+    }

tests/baselines/reference/controlFlowOptionalChain.js

@@ -588,7 +588,24 @@ while (arr[i]?.tag === "left") {
         console.log("I should ALSO be reachable");
     }
 }
-
+
+
+// Repro from #51941
+
+type Test5 = {
+  main?: {
+    childs: Record<string, Test5>;
+  };
+};
+
+function f50(obj: Test5) {
+   for (const key in obj.main?.childs) {
+      if (obj.main.childs[key] === obj) {
+        return obj;
+      }
+   }
+   return null;
+}
 
 //// [controlFlowOptionalChain.js]
 "use strict";
@@ -1091,3 +1108,12 @@ while (((_u = arr[i]) === null || _u === void 0 ? void 0 : _u.tag) === "left") {
         console.log("I should ALSO be reachable");
     }
 }
+function f50(obj) {
+    var _a;
+    for (var key in (_a = obj.main) === null || _a === void 0 ? void 0 : _a.childs) {
+        if (obj.main.childs[key] === obj) {
+            return obj;
+        }
+    }
+    return null;
+}

tests/baselines/reference/controlFlowOptionalChain.symbols

@@ -1832,3 +1832,48 @@ while (arr[i]?.tag === "left") {
     }
 }
 
+
+// Repro from #51941
+
+type Test5 = {
+>Test5 : Symbol(Test5, Decl(controlFlowOptionalChain.ts, 588, 1))
+
+  main?: {
+>main : Symbol(main, Decl(controlFlowOptionalChain.ts, 593, 14))
+
+    childs: Record<string, Test5>;
+>childs : Symbol(childs, Decl(controlFlowOptionalChain.ts, 594, 10))
+>Record : Symbol(Record, Decl(lib.es5.d.ts, --, --))
+>Test5 : Symbol(Test5, Decl(controlFlowOptionalChain.ts, 588, 1))
+
+  };
+};
+
+function f50(obj: Test5) {
+>f50 : Symbol(f50, Decl(controlFlowOptionalChain.ts, 597, 2))
+>obj : Symbol(obj, Decl(controlFlowOptionalChain.ts, 599, 13))
+>Test5 : Symbol(Test5, Decl(controlFlowOptionalChain.ts, 588, 1))
+
+   for (const key in obj.main?.childs) {
+>key : Symbol(key, Decl(controlFlowOptionalChain.ts, 600, 13))
+>obj.main?.childs : Symbol(childs, Decl(controlFlowOptionalChain.ts, 594, 10))
+>obj.main : Symbol(main, Decl(controlFlowOptionalChain.ts, 593, 14))
+>obj : Symbol(obj, Decl(controlFlowOptionalChain.ts, 599, 13))
+>main : Symbol(main, Decl(controlFlowOptionalChain.ts, 593, 14))
+>childs : Symbol(childs, Decl(controlFlowOptionalChain.ts, 594, 10))
+
+      if (obj.main.childs[key] === obj) {
+>obj.main.childs : Symbol(childs, Decl(controlFlowOptionalChain.ts, 594, 10))
+>obj.main : Symbol(main, Decl(controlFlowOptionalChain.ts, 593, 14))
+>obj : Symbol(obj, Decl(controlFlowOptionalChain.ts, 599, 13))
+>main : Symbol(main, Decl(controlFlowOptionalChain.ts, 593, 14))
+>childs : Symbol(childs, Decl(controlFlowOptionalChain.ts, 594, 10))
+>key : Symbol(key, Decl(controlFlowOptionalChain.ts, 600, 13))
+>obj : Symbol(obj, Decl(controlFlowOptionalChain.ts, 599, 13))
+
+        return obj;
+>obj : Symbol(obj, Decl(controlFlowOptionalChain.ts, 599, 13))
+      }
+   }
+   return null;
+}

tests/baselines/reference/controlFlowOptionalChain.types

@@ -2073,3 +2073,48 @@ while (arr[i]?.tag === "left") {
     }
 }
 
+
+// Repro from #51941
+
+type Test5 = {
+>Test5 : { main?: { childs: Record<string, Test5>; } | undefined; }
+
+  main?: {
+>main : { childs: Record<string, Test5>; } | undefined
+
+    childs: Record<string, Test5>;
+>childs : Record<string, Test5>
+
+  };
+};
+
+function f50(obj: Test5) {
+>f50 : (obj: Test5) => Test5 | null
+>obj : Test5
+
+   for (const key in obj.main?.childs) {
+>key : string
+>obj.main?.childs : Record<string, Test5> | undefined
+>obj.main : { childs: Record<string, Test5>; } | undefined
+>obj : Test5
+>main : { childs: Record<string, Test5>; } | undefined
+>childs : Record<string, Test5> | undefined
+
+      if (obj.main.childs[key] === obj) {
+>obj.main.childs[key] === obj : boolean
+>obj.main.childs[key] : Test5
+>obj.main.childs : Record<string, Test5>
+>obj.main : { childs: Record<string, Test5>; }
+>obj : Test5
+>main : { childs: Record<string, Test5>; }
+>childs : Record<string, Test5>
+>key : string
+>obj : Test5
+
+        return obj;
+>obj : Test5
+      }
+   }
+   return null;
+>null : null
+}

tests/cases/conformance/controlFlow/controlFlowOptionalChain.ts

@@ -590,3 +590,21 @@ while (arr[i]?.tag === "left") {
         console.log("I should ALSO be reachable");
     }
 }
+
+
+// Repro from #51941
+
+type Test5 = {
+  main?: {
+    childs: Record<string, Test5>;
+  };
+};
+
+function f50(obj: Test5) {
+   for (const key in obj.main?.childs) {
+      if (obj.main.childs[key] === obj) {
+        return obj;
+      }
+   }
+   return null;
+}