Fix FindFirstChild and WaitForChild magic functions

- Ensure FFC returns Instance? on DM types - Support second parameters for FFC and WFC Closes #704
JohnnyMorganz · Jul 20, 2024 · b3977ed · b3977ed
1 parent de93d0c
commit b3977ed
Show file tree

Hide file tree

Showing 3 changed files with 259 additions and 66 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -10,6 +10,12 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 - Sync to upstream Luau 0.635
 
+### Fixed
+
+- Fixed `:FindFirstChild()` returning `Instance` instead of `Instance?` when applied to DataModel types
+- Fixed `:FindFirstChild()` not supporting a boolean "recursive" 2nd parameter on DataModel types ([#704](https://github.com/JohnnyMorganz/luau-lsp/issues/704))
+- Fixed `:WaitForChild()` not supporting a number "timeout" 2nd parameter on DataModel types ([#704](https://github.com/JohnnyMorganz/luau-lsp/issues/704))
+
 ## [1.32.0] - 2024-07-14
 
 ### Added

diff --git a/src/platform/roblox/RobloxSourcemap.cpp b/src/platform/roblox/RobloxSourcemap.cpp
@@ -56,81 +56,58 @@ static void injectChildrenLookupFunctions(
 {
     if (auto instanceType = getTypeIdForClass(globals.globalScope, "Instance"))
     {
-        auto findFirstChildFunction = Luau::makeFunction(arena, ty, {globals.builtinTypes->stringType}, {"name"}, {*instanceType});
-        Luau::attachMagicFunction(findFirstChildFunction,
-            [node, &arena, &globals](Luau::TypeChecker& typeChecker, const Luau::ScopePtr& scope, const Luau::AstExprCall& expr,
-                const Luau::WithPredicate<Luau::TypePackId>& withPredicate) -> std::optional<Luau::WithPredicate<Luau::TypePackId>>
-            {
-                if (expr.args.size < 1)
-                    return std::nullopt;
-
-                auto str = expr.args.data[0]->as<Luau::AstExprConstantString>();
-                if (!str)
-                    return std::nullopt;
-
-                if (auto child = node->findChild(std::string(str->value.data, str->value.size)))
-                    return Luau::WithPredicate<Luau::TypePackId>{arena.addTypePack({getSourcemapType(globals, arena, *child)})};
-
-                return std::nullopt;
-            });
-        Luau::attachDcrMagicFunction(findFirstChildFunction,
-            [node, &arena, &globals](Luau::MagicFunctionCallContext context) -> bool
-            {
-                if (context.callSite->args.size < 1)
-                    return false;
-
-                auto str = context.callSite->args.data[0]->as<Luau::AstExprConstantString>();
-                if (!str)
-                    return false;
+        auto optionalInstanceType = Luau::makeOption(globals.builtinTypes, arena, *instanceType);
+        auto findFirstChildFunction = Luau::makeFunction(arena, ty,
+            {globals.builtinTypes->stringType, Luau::makeOption(globals.builtinTypes, arena, globals.builtinTypes->booleanType)},
+            {"name", "recursive"}, {optionalInstanceType});
+        auto waitForChildFunction = Luau::makeFunction(arena, ty, {globals.builtinTypes->stringType}, {"name"}, {*instanceType});
+        auto waitForChildWithTimeoutFunction = Luau::makeFunction(
+            arena, ty, {globals.builtinTypes->stringType, globals.builtinTypes->numberType}, {"name", "timeout"}, {optionalInstanceType});
 
-                if (auto child = node->findChild(std::string(str->value.data, str->value.size)))
+        for (const auto& lookupFuncTy : {findFirstChildFunction, waitForChildFunction, waitForChildWithTimeoutFunction})
+        {
+            Luau::attachMagicFunction(lookupFuncTy,
+                [node, &arena, &globals](Luau::TypeChecker& typeChecker, const Luau::ScopePtr& scope, const Luau::AstExprCall& expr,
+                    const Luau::WithPredicate<Luau::TypePackId>& withPredicate) -> std::optional<Luau::WithPredicate<Luau::TypePackId>>
                 {
-                    asMutable(context.result)
-                        ->ty.emplace<Luau::BoundTypePack>(context.solver->arena->addTypePack({getSourcemapType(globals, arena, *child)}));
-                    return true;
-                }
-                return false;
-            });
-        ctv->props["FindFirstChild"] = Luau::makeProperty(findFirstChildFunction, "@roblox/globaltype/Instance.FindFirstChild");
-        Luau::attachTag(ctv->props["FindFirstChild"].type(), kSourcemapGeneratedTag);
-        Luau::attachTag(ctv->props["FindFirstChild"].type(), "Children");
+                    if (expr.args.size < 1)
+                        return std::nullopt;
 
-        auto waitForChildFunction = Luau::makeFunction(arena, ty, {globals.builtinTypes->stringType}, {"name"}, {*instanceType});
-        Luau::attachMagicFunction(waitForChildFunction,
-            [node, &arena, &globals](Luau::TypeChecker& typeChecker, const Luau::ScopePtr& scope, const Luau::AstExprCall& expr,
-                const Luau::WithPredicate<Luau::TypePackId>& withPredicate) -> std::optional<Luau::WithPredicate<Luau::TypePackId>>
-            {
-                if (expr.args.size < 1)
-                    return std::nullopt;
+                    auto str = expr.args.data[0]->as<Luau::AstExprConstantString>();
+                    if (!str)
+                        return std::nullopt;
+
+                    if (auto child = node->findChild(std::string(str->value.data, str->value.size)))
+                        return Luau::WithPredicate<Luau::TypePackId>{arena.addTypePack({getSourcemapType(globals, arena, *child)})};
 
-                auto str = expr.args.data[0]->as<Luau::AstExprConstantString>();
-                if (!str)
                     return std::nullopt;
+                });
+            Luau::attachDcrMagicFunction(lookupFuncTy,
+                [node, &arena, &globals](Luau::MagicFunctionCallContext context) -> bool
+                {
+                    if (context.callSite->args.size < 1)
+                        return false;
 
-                if (auto child = node->findChild(std::string(str->value.data, str->value.size)))
-                    return Luau::WithPredicate<Luau::TypePackId>{arena.addTypePack({getSourcemapType(globals, arena, *child)})};
+                    auto str = context.callSite->args.data[0]->as<Luau::AstExprConstantString>();
+                    if (!str)
+                        return false;
 
-                return std::nullopt;
-            });
-        Luau::attachDcrMagicFunction(waitForChildFunction,
-            [node, &arena, &globals](Luau::MagicFunctionCallContext context) -> bool
-            {
-                if (context.callSite->args.size < 1)
+                    if (auto child = node->findChild(std::string(str->value.data, str->value.size)))
+                    {
+                        asMutable(context.result)
+                            ->ty.emplace<Luau::BoundTypePack>(context.solver->arena->addTypePack({getSourcemapType(globals, arena, *child)}));
+                        return true;
+                    }
                     return false;
+                });
+        }
 
-                auto str = context.callSite->args.data[0]->as<Luau::AstExprConstantString>();
-                if (!str)
-                    return false;
+        ctv->props["FindFirstChild"] = Luau::makeProperty(findFirstChildFunction, "@roblox/globaltype/Instance.FindFirstChild");
+        Luau::attachTag(ctv->props["FindFirstChild"].type(), kSourcemapGeneratedTag);
+        Luau::attachTag(ctv->props["FindFirstChild"].type(), "Children");
 
-                if (auto child = node->findChild(std::string(str->value.data, str->value.size)))
-                {
-                    asMutable(context.result)
-                        ->ty.emplace<Luau::BoundTypePack>(context.solver->arena->addTypePack({getSourcemapType(globals, arena, *child)}));
-                    return true;
-                }
-                return false;
-            });
-        ctv->props["WaitForChild"] = Luau::makeProperty(waitForChildFunction, "@roblox/globaltype/Instance.WaitForChild");
+        ctv->props["WaitForChild"] = Luau::makeProperty(
+            Luau::makeIntersection(arena, {waitForChildFunction, waitForChildWithTimeoutFunction}), "@roblox/globaltype/Instance.WaitForChild");
         Luau::attachTag(ctv->props["WaitForChild"].type(), kSourcemapGeneratedTag);
         Luau::attachTag(ctv->props["WaitForChild"].type(), "Children");
     }
@@ -185,7 +162,8 @@ static Luau::TypeId getSourcemapType(const Luau::GlobalTypes& globals, Luau::Typ
 
             // Create the ClassType representing the instance
             std::string typeName = types::getTypeName(baseTypeId).value_or(node->name);
-            Luau::ClassType baseInstanceCtv{typeName, {}, baseTypeId, instanceMetaIdentity, {}, {}, instanceCtv->definitionModuleName, instanceCtv->definitionLocation};
+            Luau::ClassType baseInstanceCtv{
+                typeName, {}, baseTypeId, instanceMetaIdentity, {}, {}, instanceCtv->definitionModuleName, instanceCtv->definitionLocation};
             auto typeId = arena.addType(std::move(baseInstanceCtv));
 
             // Attach Parent and Children info

diff --git a/tests/Sourcemap.test.cpp b/tests/Sourcemap.test.cpp
@@ -96,6 +96,118 @@ TEST_CASE_FIXTURE(Fixture, "can_access_children_via_find_first_child")
     CHECK(Luau::toString(requireType("head")) == "Part");
 }
 
+TEST_CASE_FIXTURE(Fixture, "find_first_child_handles_unknown_child")
+{
+    client->globalConfig.diagnostics.strictDatamodelTypes = true;
+    loadSourcemap(R"(
+            {
+                "name": "Game",
+                "className": "DataModel",
+                "children": [
+                    {
+                        "name": "TemplateR15",
+                        "className": "Part",
+                        "children": [
+                            {"name": "Head", "className": "Part"}
+                        ]
+                    }
+                ]
+            }
+        )");
+
+    auto result = check(R"(
+        --!strict
+        local template = game:FindFirstChild("Unknown")
+    )");
+
+    LUAU_LSP_REQUIRE_NO_ERRORS(result);
+    CHECK(Luau::toString(requireType("template")) == "Instance?");
+}
+
+TEST_CASE_FIXTURE(Fixture, "find_first_child_works_without_sourcemap")
+{
+    client->globalConfig.diagnostics.strictDatamodelTypes = true;
+
+    auto result = check(R"(
+        --!strict
+        local template = game:FindFirstChild("Unknown")
+    )");
+
+    LUAU_LSP_REQUIRE_NO_ERRORS(result);
+    CHECK(Luau::toString(requireType("template")) == "Instance?");
+}
+
+
+TEST_CASE_FIXTURE(Fixture, "find_first_child_still_supports_recursive_parameter_with_sourcemap")
+{
+    client->globalConfig.diagnostics.strictDatamodelTypes = true;
+    loadSourcemap(R"(
+            {
+                "name": "Game",
+                "className": "DataModel",
+                "children": [
+                    {
+                        "name": "TemplateR15",
+                        "className": "Part",
+                        "children": [
+                            {"name": "Head", "className": "Part"}
+                        ]
+                    }
+                ]
+            }
+        )");
+
+    // TODO: Support recursive datamodel lookup - https://github.com/JohnnyMorganz/luau-lsp/issues/689
+    auto result = check(R"(
+        --!strict
+        local template = game:FindFirstChild("Head", true)
+    )");
+
+    LUAU_LSP_REQUIRE_NO_ERRORS(result);
+    CHECK(Luau::toString(requireType("template")) == "Instance?");
+}
+
+TEST_CASE_FIXTURE(Fixture, "find_first_child_still_supports_recursive_parameter_without_sourcemap")
+{
+    client->globalConfig.diagnostics.strictDatamodelTypes = true;
+    auto result = check(R"(
+        --!strict
+        local template = game:FindFirstChild("Unknown", true)
+    )");
+
+    LUAU_LSP_REQUIRE_NO_ERRORS(result);
+    CHECK(Luau::toString(requireType("template")) == "Instance?");
+}
+
+
+TEST_CASE_FIXTURE(Fixture, "find_first_child_finds_direct_child_recursively")
+{
+    client->globalConfig.diagnostics.strictDatamodelTypes = true;
+    loadSourcemap(R"(
+            {
+                "name": "Game",
+                "className": "DataModel",
+                "children": [
+                    {
+                        "name": "TemplateR15",
+                        "className": "Part",
+                        "children": [
+                            {"name": "Head", "className": "Part"}
+                        ]
+                    }
+                ]
+            }
+        )");
+
+    auto result = check(R"(
+        --!strict
+        local template = game:FindFirstChild("TemplateR15", true)
+    )");
+
+    LUAU_LSP_REQUIRE_NO_ERRORS(result);
+    CHECK(Luau::toString(requireType("template")) == "Part");
+}
+
 TEST_CASE_FIXTURE(Fixture, "can_access_children_via_wait_for_child")
 {
     client->globalConfig.diagnostics.strictDatamodelTypes = true;
@@ -126,6 +238,103 @@ TEST_CASE_FIXTURE(Fixture, "can_access_children_via_wait_for_child")
     CHECK(Luau::toString(requireType("head")) == "Part");
 }
 
+TEST_CASE_FIXTURE(Fixture, "wait_for_child_handles_unknown_child")
+{
+    client->globalConfig.diagnostics.strictDatamodelTypes = true;
+    loadSourcemap(R"(
+            {
+                "name": "Game",
+                "className": "DataModel",
+                "children": [
+                    {
+                        "name": "TemplateR15",
+                        "className": "Part",
+                        "children": [
+                            {"name": "Head", "className": "Part"}
+                        ]
+                    }
+                ]
+            }
+        )");
+
+    auto result = check(R"(
+        --!strict
+        local template = game:WaitForChild("UnknownChild")
+    )");
+
+    LUAU_LSP_REQUIRE_NO_ERRORS(result);
+    CHECK(Luau::toString(requireType("template")) == "Instance");
+}
+
+TEST_CASE_FIXTURE(Fixture, "wait_for_child_still_supports_timeout_parameter_with_sourcemap")
+{
+    client->globalConfig.diagnostics.strictDatamodelTypes = true;
+    loadSourcemap(R"(
+            {
+                "name": "Game",
+                "className": "DataModel",
+                "children": [
+                    {
+                        "name": "TemplateR15",
+                        "className": "Part",
+                        "children": [
+                            {"name": "Head", "className": "Part"}
+                        ]
+                    }
+                ]
+            }
+        )");
+
+    auto result = check(R"(
+        --!strict
+        local template = game:WaitForChild("UnknownChild", 10)
+    )");
+
+    LUAU_LSP_REQUIRE_NO_ERRORS(result);
+    CHECK(Luau::toString(requireType("template")) == "Instance?");
+}
+
+TEST_CASE_FIXTURE(Fixture, "wait_for_child_still_supports_timeout_parameter_without_sourcemap")
+{
+    client->globalConfig.diagnostics.strictDatamodelTypes = true;
+
+    auto result = check(R"(
+        --!strict
+        local template = game:WaitForChild("TemplateR15", 10)
+    )");
+
+    LUAU_LSP_REQUIRE_NO_ERRORS(result);
+    CHECK(Luau::toString(requireType("template")) == "Instance?");
+}
+
+TEST_CASE_FIXTURE(Fixture, "wait_for_child_finds_direct_child_with_timeout_parameter")
+{
+    client->globalConfig.diagnostics.strictDatamodelTypes = true;
+    loadSourcemap(R"(
+            {
+                "name": "Game",
+                "className": "DataModel",
+                "children": [
+                    {
+                        "name": "TemplateR15",
+                        "className": "Part",
+                        "children": [
+                            {"name": "Head", "className": "Part"}
+                        ]
+                    }
+                ]
+            }
+        )");
+
+    auto result = check(R"(
+        --!strict
+        local template = game:WaitForChild("TemplateR15", 10)
+    )");
+
+    LUAU_LSP_REQUIRE_NO_ERRORS(result);
+    CHECK(Luau::toString(requireType("template")) == "Part");
+}
+
 TEST_CASE_FIXTURE(Fixture, "relative_and_absolute_types_are_consistent")
 {
     client->globalConfig.diagnostics.strictDatamodelTypes = true;