Add instance type resolution for FindFirstChild/FindFirstAncestor

Closes #18

Author: JohnnyMorganz

CHANGELOG.md

@@ -9,6 +9,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 ### Added
 
 - Player will now have `PlayerGui`, `Backpack`, `StarterGear` and `PlayerScripts` children, with the relevant Starter instances copied into it (StarterGui, StarterPack, PlayerScripts)
+- `Instance:FindFirstChild("name")` and `Instance:FindFirstAncestor("name")` will now correctly resolve to the relevant instance type if found. This allows the type checker to correctly resolve children/parents etc.
 
 ### Changed
 

src/LuauExt.cpp

@@ -78,6 +78,55 @@ Luau::TypeId makeLazyInstanceType(Luau::TypeArena& arena, const Luau::ScopePtr&
             {
                 ctv->props[child->name] = Luau::makeProperty(makeLazyInstanceType(arena, globalScope, child, typeId));
             }
+
+            // Add FindFirstAncestor and FindFirstChild
+            if (auto instanceType = getTypeIdForClass(globalScope, "Instance"))
+            {
+                auto findFirstAncestorFunction = Luau::makeFunction(arena, typeId, {Luau::getSingletonTypes().stringType}, {"name"}, {*instanceType});
+                Luau::attachMagicFunction(findFirstAncestorFunction,
+                    [node](Luau::TypeChecker& typeChecker, const Luau::ScopePtr& scope, const Luau::AstExprCall& expr,
+                        Luau::ExprResult<Luau::TypePackId> exprResult) -> std::optional<Luau::ExprResult<Luau::TypePackId>>
+                    {
+                        if (expr.args.size < 1)
+                            return std::nullopt;
+
+                        auto str = expr.args.data[0]->as<Luau::AstExprConstantString>();
+                        if (!str)
+                            return std::nullopt;
+
+                        // This is a O(n) search, not great!
+                        if (auto ancestor = node->findAncestor(std::string(str->value.data, str->value.size)))
+                        {
+                            return Luau::ExprResult<Luau::TypePackId>{
+                                typeChecker.globalTypes.addTypePack({makeLazyInstanceType(typeChecker.globalTypes, scope, *ancestor, std::nullopt)})};
+                        }
+
+                        return std::nullopt;
+                    });
+                ctv->props["FindFirstAncestor"] = Luau::makeProperty(findFirstAncestorFunction, "@roblox/globaltype/Instance.FindFirstAncestor");
+
+                auto findFirstChildFunction = Luau::makeFunction(arena, typeId, {Luau::getSingletonTypes().stringType}, {"name"}, {*instanceType});
+                Luau::attachMagicFunction(findFirstChildFunction,
+                    [node, typeId](Luau::TypeChecker& typeChecker, const Luau::ScopePtr& scope, const Luau::AstExprCall& expr,
+                        Luau::ExprResult<Luau::TypePackId> exprResult) -> std::optional<Luau::ExprResult<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::ExprResult<Luau::TypePackId>{
+                                typeChecker.globalTypes.addTypePack({makeLazyInstanceType(typeChecker.globalTypes, scope, *child, typeId)})};
+                        }
+
+                        return std::nullopt;
+                    });
+                ctv->props["FindFirstChild"] = Luau::makeProperty(findFirstChildFunction, "@roblox/globaltype/Instance.FindFirstChild");
+            }
         }
         return typeId;
     };

src/Sourcemap.cpp

@@ -55,6 +55,18 @@ std::optional<SourceNodePtr> SourceNode::findChild(const std::string& name)
     return std::nullopt;
 }
 
+std::optional<SourceNodePtr> SourceNode::findAncestor(const std::string& name)
+{
+    auto current = parent;
+    while (auto currentPtr = current.lock())
+    {
+        if (currentPtr->name == name)
+            return currentPtr;
+        current = currentPtr->parent;
+    }
+    return std::nullopt;
+}
+
 static bool endsWith(std::string str, std::string suffix)
 {
     return str.size() >= suffix.size() && 0 == str.compare(str.size() - suffix.size(), suffix.size(), suffix);

src/include/LSP/Sourcemap.hpp

@@ -20,6 +20,8 @@ struct SourceNode
     std::optional<std::filesystem::path> getScriptFilePath();
     Luau::SourceCode::Type sourceCodeType() const;
     std::optional<SourceNodePtr> findChild(const std::string& name);
+    // O(n) search for ancestor of name
+    std::optional<SourceNodePtr> findAncestor(const std::string& name);
 };
 
 static void from_json(const json& j, SourceNode& p)