From 97db8166a51a138d942542f81ca5e2cb01d89392 Mon Sep 17 00:00:00 2001
From: Kampfkarren <3190756+Kampfkarren@users.noreply.github.com>
Date: Thu, 28 Nov 2024 22:34:35 -0800
Subject: [PATCH] Fix values preserving nil in iterators

---
 Analysis/src/TypeInfer.cpp     | 14 +++++++++++++-
 tests/TypeInfer.loops.test.cpp | 31 +++++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/Analysis/src/TypeInfer.cpp b/Analysis/src/TypeInfer.cpp
index 7ed3290bb..b093daba1 100644
--- a/Analysis/src/TypeInfer.cpp
+++ b/Analysis/src/TypeInfer.cpp
@@ -34,6 +34,7 @@ LUAU_FASTFLAGVARIABLE(DebugLuauFreezeDuringUnification)
 LUAU_FASTFLAG(LuauInstantiateInSubtyping)
 LUAU_FASTFLAGVARIABLE(LuauMetatableFollow)
 LUAU_FASTFLAGVARIABLE(LuauRequireCyclesDontAlwaysReturnAny)
+LUAU_FASTFLAGVARIABLE(LuauNilInForLoops)
 
 namespace Luau
 {
@@ -1285,7 +1286,18 @@ ControlFlow TypeChecker::check(const ScopePtr& scope, const AstStatForIn& forin)
                 unify(iterTable->indexer->indexType, varTypes[0], scope, forin.location);
 
             if (varTypes.size() > 1)
-                unify(iterTable->indexer->indexResultType, varTypes[1], scope, forin.location);
+            {
+                if (FFlag::LuauNilInForLoops)
+                {
+                    std::optional<TypeId> withoutNilTy = tryStripUnionFromNil(iterTable->indexer->indexResultType);
+
+                    unify(withoutNilTy ? *withoutNilTy : iterTable->indexer->indexResultType, varTypes[1], scope, forin.location);
+                }
+                else
+                {
+                    unify(iterTable->indexer->indexResultType, varTypes[1], scope, forin.location);
+                }
+            }
 
             for (size_t i = 2; i < varTypes.size(); ++i)
                 unify(nilType, varTypes[i], scope, forin.location);
diff --git a/tests/TypeInfer.loops.test.cpp b/tests/TypeInfer.loops.test.cpp
index 41753b66d..fa689afe3 100644
--- a/tests/TypeInfer.loops.test.cpp
+++ b/tests/TypeInfer.loops.test.cpp
@@ -15,6 +15,7 @@
 using namespace Luau;
 
 LUAU_FASTFLAG(LuauSolverV2)
+LUAU_FASTFLAG(LuauNilInForLoops)
 
 TEST_SUITE_BEGIN("TypeInferLoops");
 
@@ -1255,4 +1256,34 @@ end
     )");
 }
 
+TEST_CASE_FIXTURE(Fixture, "nil_in_for_loops")
+{
+    ScopedFastFlag sff = {FFlag::LuauNilInForLoops, true};
+
+    LUAU_REQUIRE_NO_ERRORS(check(R"(
+        --!strict
+        local function f(x: { [string]: string? })
+            for key, value in x do
+                local v: string = value
+            end
+        end
+    )"));
+}
+
+// Custom iterators can return nil, just not without
+TEST_CASE_FIXTURE(BuiltinsFixture, "nil_in_for_loops_iter")
+{
+    LUAU_REQUIRE_NO_ERRORS(check(R"(
+        --!strict
+        type T = typeof(setmetatable({}, {} :: {
+            __iter: (any) -> () -> (boolean, string?)
+        }))
+
+        local function f(x: T)
+            for a: boolean, b: string? in x do
+            end
+        end
+    )"));
+}
+
 TEST_SUITE_END();