From 4d98a16ea8f4985ca757d2ceebc278a6db3f9761 Mon Sep 17 00:00:00 2001
From: JohnnyMorganz <johnnymorganz@outlook.com>
Date: Tue, 5 Jul 2022 17:08:05 +0100
Subject: [PATCH] Store resolved types in `astResolvedTypes` (#574)

---
 Analysis/include/Luau/TypeInfer.h |  1 +
 Analysis/src/Frontend.cpp         |  2 ++
 Analysis/src/TypeInfer.cpp        | 27 ++++++++++++++++++++-------
 tests/Frontend.test.cpp           |  2 ++
 tests/TypeInfer.test.cpp          | 23 +++++++++++++++++++++++
 5 files changed, 48 insertions(+), 7 deletions(-)

diff --git a/Analysis/include/Luau/TypeInfer.h b/Analysis/include/Luau/TypeInfer.h
index 455654d..dac2902 100644
--- a/Analysis/include/Luau/TypeInfer.h
+++ b/Analysis/include/Luau/TypeInfer.h
@@ -345,6 +345,7 @@ private:
     TypePackId freshTypePack(TypeLevel level);
 
     TypeId resolveType(const ScopePtr& scope, const AstType& annotation);
+    TypeId resolveTypeWorker(const ScopePtr& scope, const AstType& annotation);
     TypePackId resolveTypePack(const ScopePtr& scope, const AstTypeList& types);
     TypePackId resolveTypePack(const ScopePtr& scope, const AstTypePack& annotation);
     TypeId instantiateTypeFun(const ScopePtr& scope, const TypeFun& tf, const std::vector<TypeId>& typeParams,
diff --git a/Analysis/src/Frontend.cpp b/Analysis/src/Frontend.cpp
index 4cfaa11..9195363 100644
--- a/Analysis/src/Frontend.cpp
+++ b/Analysis/src/Frontend.cpp
@@ -496,6 +496,8 @@ CheckResult Frontend::check(const ModuleName& name, std::optional<FrontendOption
             module->astTypes.clear();
             module->astExpectedTypes.clear();
             module->astOriginalCallTypes.clear();
+            module->astResolvedTypes.clear();
+            module->astResolvedTypePacks.clear();
             module->scopes.resize(1);
         }
 
diff --git a/Analysis/src/TypeInfer.cpp b/Analysis/src/TypeInfer.cpp
index d9486a4..4fafb50 100644
--- a/Analysis/src/TypeInfer.cpp
+++ b/Analysis/src/TypeInfer.cpp
@@ -4884,6 +4884,13 @@ TypePackId TypeChecker::freshTypePack(TypeLevel level)
 }
 
 TypeId TypeChecker::resolveType(const ScopePtr& scope, const AstType& annotation)
+{
+    TypeId ty = resolveTypeWorker(scope, annotation);
+    currentModule->astResolvedTypes[&annotation] = ty;
+    return ty;
+}
+
+TypeId TypeChecker::resolveTypeWorker(const ScopePtr& scope, const AstType& annotation)
 {
     if (const auto& lit = annotation.as<AstTypeReference>())
     {
@@ -5200,9 +5207,10 @@ TypePackId TypeChecker::resolveTypePack(const ScopePtr& scope, const AstTypeList
 
 TypePackId TypeChecker::resolveTypePack(const ScopePtr& scope, const AstTypePack& annotation)
 {
+    TypePackId result;
     if (const AstTypePackVariadic* variadic = annotation.as<AstTypePackVariadic>())
     {
-        return addTypePack(TypePackVar{VariadicTypePack{resolveType(scope, *variadic->variadicType)}});
+        result = addTypePack(TypePackVar{VariadicTypePack{resolveType(scope, *variadic->variadicType)}});
     }
     else if (const AstTypePackGeneric* generic = annotation.as<AstTypePackGeneric>())
     {
@@ -5216,10 +5224,12 @@ TypePackId TypeChecker::resolveTypePack(const ScopePtr& scope, const AstTypePack
             else
                 reportError(TypeError{generic->location, UnknownSymbol{genericName, UnknownSymbol::Type}});
 
-            return errorRecoveryTypePack(scope);
+            result = errorRecoveryTypePack(scope);
+        }
+        else
+        {
+            result = *genericTy;
         }
-
-        return *genericTy;
     }
     else if (const AstTypePackExplicit* explicitTp = annotation.as<AstTypePackExplicit>())
     {
@@ -5229,14 +5239,17 @@ TypePackId TypeChecker::resolveTypePack(const ScopePtr& scope, const AstTypePack
             types.push_back(resolveType(scope, *type));
 
         if (auto tailType = explicitTp->typeList.tailType)
-            return addTypePack(types, resolveTypePack(scope, *tailType));
-
-        return addTypePack(types);
+            result = addTypePack(types, resolveTypePack(scope, *tailType));
+        else
+            result = addTypePack(types);
     }
     else
     {
         ice("Unknown AstTypePack kind");
     }
+
+    currentModule->astResolvedTypePacks[&annotation] = result;
+    return result;
 }
 
 bool ApplyTypeFunction::isDirty(TypeId ty)
diff --git a/tests/Frontend.test.cpp b/tests/Frontend.test.cpp
index b9c2470..4f33139 100644
--- a/tests/Frontend.test.cpp
+++ b/tests/Frontend.test.cpp
@@ -791,6 +791,8 @@ TEST_CASE_FIXTURE(FrontendFixture, "discard_type_graphs")
     CHECK_EQ(0, module->internalTypes.typeVars.size());
     CHECK_EQ(0, module->internalTypes.typePacks.size());
     CHECK_EQ(0, module->astTypes.size());
+    CHECK_EQ(0, module->astResolvedTypes.size());
+    CHECK_EQ(0, module->astResolvedTypePacks.size());
 }
 
 TEST_CASE_FIXTURE(FrontendFixture, "it_should_be_safe_to_stringify_errors_when_full_type_graph_is_discarded")
diff --git a/tests/TypeInfer.test.cpp b/tests/TypeInfer.test.cpp
index efdfe0b..a175b82 100644
--- a/tests/TypeInfer.test.cpp
+++ b/tests/TypeInfer.test.cpp
@@ -1003,4 +1003,27 @@ TEST_CASE_FIXTURE(Fixture, "do_not_bind_a_free_table_to_a_union_containing_that_
     )");
 }
 
+TEST_CASE_FIXTURE(Fixture, "types stored in astResolvedTypes")
+{
+    CheckResult result = check(R"(
+type alias = typeof("hello")
+local function foo(param: alias)
+end
+    )");
+
+    auto node = findNodeAtPosition(*getMainSourceModule(), {2, 16});
+    auto ty = lookupType("alias");
+    REQUIRE(node);
+    REQUIRE(node->is<AstExprFunction>());
+    REQUIRE(ty);
+
+    auto func = node->as<AstExprFunction>();
+    REQUIRE(func->args.size == 1);
+
+    auto arg = *func->args.begin();
+    auto annotation = arg->annotation;
+
+    CHECK_EQ(*getMainModule()->astResolvedTypes.find(annotation), *ty);
+}
+
 TEST_SUITE_END();