Mark explicit self functions as methods for semantic tokens

Also creates an `isMethod` helper

Closes #531

CHANGELOG.md

@@ -14,6 +14,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 - Switched to memory-efficient implementation of workspace diagnostics (currently behind FFlag `LuauStacklessTypeClone3`)
 - Improved handling of configuration info received from non-VSCode clients
+- Functions with explicitly defined `self` parameters are correctly marked with the `method` semantic token
 
 ## [1.27.0] - 2023-12-25
 

src/LuauExt.cpp

@@ -1296,3 +1296,17 @@ bool isRequire(const Luau::AstExpr* expr)
 
     return false;
 }
+
+bool isMethod(const Luau::FunctionType* ftv)
+{
+    if (ftv->hasSelf)
+        return true;
+
+    // TODO: hasSelf is not always specified, so we manually check for the "self" name (https://github.com/Roblox/luau/issues/551)
+    if (ftv->argNames.size() > 0 && ftv->argNames[0].has_value() && ftv->argNames[0]->name == "self")
+    {
+        return true;
+    }
+
+    return false;
+}

src/include/LSP/LuauExt.hpp

@@ -74,6 +74,7 @@ lsp::Diagnostic createParseErrorDiagnostic(const Luau::ParseError& error, const
 
 bool isGetService(const Luau::AstExpr* expr);
 bool isRequire(const Luau::AstExpr* expr);
+bool isMethod(const Luau::FunctionType* ftv);
 
 struct FindImportsVisitor : public Luau::AstVisitor
 {

src/include/LSP/SemanticTokens.hpp

@@ -1,6 +1,7 @@
 #pragma once
-#include "Luau/Ast.h"
+#include "Luau/Location.h"
 #include "Luau/Module.h"
+#include "Luau/Frontend.h"
 #include "Protocol/SemanticTokens.hpp"
 
 struct SemanticToken
@@ -11,4 +12,4 @@ struct SemanticToken
     lsp::SemanticTokenModifiers tokenModifiers;
 };
 
-std::vector<SemanticToken> getSemanticTokens(const Luau::ModulePtr& module, Luau::SourceModule* sourceModule);
+std::vector<SemanticToken> getSemanticTokens(const Luau::Frontend& frontend, const Luau::ModulePtr& module, const Luau::SourceModule* sourceModule);

src/operations/InlayHints.cpp

@@ -170,7 +170,7 @@ struct InlayHintVisitor : public Luau::AstVisitor
                     hints.emplace_back(hint);
                 }
             }
-            
+
             // Parameter types hint
             if (config.inlayHints.parameterTypes)
             {
@@ -179,8 +179,7 @@ struct InlayHintVisitor : public Luau::AstVisitor
                     return true;
 
                 // Skip first item if it is self
-                // TODO: hasSelf is not always specified, so we manually check for the "self" name (https://github.com/Roblox/luau/issues/551)
-                if (ftv->hasSelf || (func->args.size > 0 && func->args.data[0]->name == "self"))
+                if (isMethod(ftv))
                     it++;
 
                 for (auto param : func->args)
@@ -228,8 +227,7 @@ struct InlayHintVisitor : public Luau::AstVisitor
             {
                 // Skip first item if it is self
                 // TODO: hasSelf is not always specified, so we manually check for the "self" name (https://github.com/Roblox/luau/issues/551)
-                if (idx == 0 && (ftv->hasSelf || (namesIt != ftv->argNames.end() && namesIt->has_value() && namesIt->value().name == "self")) &&
-                    call->self)
+                if (idx == 0 && isMethod(ftv) && call->self)
                     namesIt++;
 
                 if (namesIt == ftv->argNames.end())

src/operations/SemanticTokens.cpp

@@ -42,7 +42,7 @@ static lsp::SemanticTokenTypes inferTokenType(const Luau::TypeId ty, lsp::Semant
 
     if (auto ftv = Luau::get<Luau::FunctionType>(followedTy))
     {
-        if (ftv->hasSelf)
+        if (isMethod(ftv))
             return lsp::SemanticTokenTypes::Method;
         else
             return lsp::SemanticTokenTypes::Function;
@@ -420,4 +420,4 @@ std::optional<lsp::SemanticTokens> LanguageServer::semanticTokens(const lsp::Sem
 {
     auto workspace = findWorkspace(params.textDocument.uri);
     return workspace->semanticTokens(params);
-}
+}

src/operations/SignatureHelp.cpp

@@ -125,12 +125,9 @@ std::optional<lsp::SignatureHelp> WorkspaceFolder::signatureHelp(const lsp::Sign
         for (; it != Luau::end(ftv->argTypes); it++, idx++)
         {
             // If the function has self, and the caller has called as a method (i.e., :), then omit the self parameter
-            // TODO: hasSelf is not always specified, so we manually check for the "self" name (https://github.com/Roblox/luau/issues/551)
-            if (idx == 0 && (ftv->hasSelf || (ftv->argNames.size() > 0 && ftv->argNames[0].has_value() && ftv->argNames[0]->name == "self")) &&
-                candidate->self)
+            if (idx == 0 && isMethod(ftv) && candidate->self)
                 continue;
-
-
+
             // Show parameter documentation
             // TODO: parse moonwave docs for param documentation?
             lsp::MarkupContent parameterDocumentation{lsp::MarkupKind::Markdown, ""};

tests/SemanticTokens.test.cpp

@@ -1,19 +1,34 @@
 #include "doctest.h"
+#include "Fixture.h"
 #include "LSP/SemanticTokens.hpp"
+
 TEST_SUITE_BEGIN("SemanticTokens");
 
-// TEST_CASE_FIXTURE(Fixture, "function_definition_name")
-// {
-//     check(R"(
-//         function foo() end
-//     )");
+std::optional<SemanticToken> getSemanticToken(const std::vector<SemanticToken>& tokens, const Luau::Position& start)
+{
+    for (const auto& token : tokens)
+        if (token.start == start)
+            return token;
+    return std::nullopt;
+}
+
+TEST_CASE_FIXTURE(Fixture, "explicit_self_method_has_method_semantic_token")
+{
+    check(R"(
+        type myClass = {
+            foo: (self: myClass) -> ()
+        }
+        local a: myClass = nil
+        a:foo()
+    )");
 
-//     auto tokens = getSemanticTokens(getMainModule(), getMainSourceModule());
-//     REQUIRE(!tokens.empty());
+    auto tokens = getSemanticTokens(workspace.frontend, getMainModule(), getMainSourceModule());
+    REQUIRE(!tokens.empty());
 
-//     auto token = *tokens.begin();
-//     CHECK_EQ(token.tokenType, lsp::SemanticTokenTypes::Function);
-//     CHECK_EQ(token.tokenModifiers, lsp::SemanticTokenModifiers::None);
-// }
+    auto token = getSemanticToken(tokens, Luau::Position{5, 10});
+    REQUIRE(token);
+    CHECK_EQ(token->tokenType, lsp::SemanticTokenTypes::Method);
+    CHECK_EQ(token->tokenModifiers, lsp::SemanticTokenModifiers::None);
+}
 
-TEST_SUITE_END();
+TEST_SUITE_END();