From 183a493625c1945d26375e9eea7197784899de87 Mon Sep 17 00:00:00 2001
From: JohnnyMorganz <johnnymorganz@outlook.com>
Date: Sat, 30 Apr 2022 21:32:38 +0100
Subject: [PATCH] Type FindFirstChildWhichIsA(T) etc. as returning T?

Done through magic functions
Closes #7
---
 main.cpp | 26 ++++++++++++++++++++++++--
 1 file changed, 24 insertions(+), 2 deletions(-)

diff --git a/main.cpp b/main.cpp
index a99d2a6..1dfdb4c 100644
--- a/main.cpp
+++ b/main.cpp
@@ -536,12 +536,29 @@ std::optional<Luau::ExprResult<Luau::TypePackId>> magicFunctionInstanceIsA(
     if (!lvalue || !tfun)
         return std::nullopt;
 
-    unfreeze(typeChecker.globalTypes);
     Luau::TypePackId booleanPack = typeChecker.globalTypes.addTypePack({typeChecker.booleanType});
-    freeze(typeChecker.globalTypes);
     return Luau::ExprResult<Luau::TypePackId>{booleanPack, {Luau::IsAPredicate{std::move(*lvalue), expr.location, tfun->type}}};
 }
 
+// Magic function for `Instance:FindFirstChildWhichIsA("ClassName")` and friends
+std::optional<Luau::ExprResult<Luau::TypePackId>> magicFunctionFindFirstXWhichIsA(
+    Luau::TypeChecker& typeChecker, const Luau::ScopePtr& scope, const Luau::AstExprCall& expr, Luau::ExprResult<Luau::TypePackId> exprResult)
+{
+    if (expr.args.size < 1)
+        return std::nullopt;
+
+    auto str = expr.args.data[0]->as<Luau::AstExprConstantString>();
+    if (!str)
+        return std::nullopt;
+
+    std::optional<Luau::TypeFun> tfun = scope->lookupType(std::string(str->value.data, str->value.size));
+    if (!tfun)
+        return std::nullopt;
+
+    Luau::TypeId nillableClass = Luau::makeOption(typeChecker, typeChecker.globalTypes, tfun->type);
+    return Luau::ExprResult<Luau::TypePackId>{typeChecker.globalTypes.addTypePack({nillableClass})};
+}
+
 int main(int argc, char** argv)
 {
     Luau::assertHandler() = assertionHandler;
@@ -815,12 +832,17 @@ int main(int argc, char** argv)
             }
 
             // Register Instance:IsA("ClassName") type predicate
+            // Register FindFirstChildWhichIsA / FindFirstChildOfClassName / FindFirstAncestorWhichIsA / FindFirstAncestorOfClass magic functions
             auto instanceType = frontend.typeChecker.globalScope->lookupType("Instance");
             if (instanceType.has_value())
             {
                 if (Luau::ClassTypeVar* ctv = Luau::getMutable<Luau::ClassTypeVar>(instanceType.value().type))
                 {
                     Luau::attachMagicFunction(ctv->props["IsA"].type, magicFunctionInstanceIsA);
+                    Luau::attachMagicFunction(ctv->props["FindFirstChildWhichIsA"].type, magicFunctionFindFirstXWhichIsA);
+                    Luau::attachMagicFunction(ctv->props["FindFirstChildOfClass"].type, magicFunctionFindFirstXWhichIsA);
+                    Luau::attachMagicFunction(ctv->props["FindFirstAncestorWhichIsA"].type, magicFunctionFindFirstXWhichIsA);
+                    Luau::attachMagicFunction(ctv->props["FindFirstAncestorOfClass"].type, magicFunctionFindFirstXWhichIsA);
                 }
             }
         }