Add support for auto imports (#86)

* Support auto service importing

* Update changelog

* Alphabeticalise location of service

* Update changelog

* Change to off by default

Author: JohnnyMorganz

CHANGELOG.md

@@ -17,6 +17,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 - Added configuration option `luau-lsp.hover.showTableKinds` (default: off) to indicate whether kinds (`{+ ... +}`, `{| ... |}`) are shown in hover information
 - Added configuration option `luau-lsp.hover.multilineFunctionDefinitions` (default: off) to spread function definitions in hover panel across multiple lines
 - Added configuration option `luau-lsp.hover.strictDatamodelTypes` (default: on) to use strict DataModel type information in hover panel (equivalent to autocomplete). When disabled, the same type information that the diagnostic type checker uses is displayed
+- Added support for automatic service importing. When using a service which has not yet been defined, it will be added (alphabetically) to the top of the file. Config setting: `luau-lsp.completion.suggestImports`
 
 ### Changed
 

editors/code/package.json

@@ -86,6 +86,11 @@
                     "type": "boolean",
                     "default": false
                 },
+                "luau-lsp.completion.suggestImports": {
+                    "markdownDescription": "Suggest automatic imports in completion items",
+                    "type": "boolean",
+                    "default": false
+                },
                 "luau-lsp.ignoreGlobs": {
                     "markdownDescription": "Diagnostics will not be reported for any file matching these globs unless the file is currently open",
                     "type": "array",

src/include/LSP/ClientConfiguration.hpp

@@ -69,9 +69,11 @@ NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE_WITH_DEFAULT(
 struct ClientCompletionConfiguration
 {
     bool enabled = true;
+    /// Whether we should suggest automatic imports in completions
+    bool suggestImports = false;
 };
 
-NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE_WITH_DEFAULT(ClientCompletionConfiguration, enabled);
+NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE_WITH_DEFAULT(ClientCompletionConfiguration, enabled, suggestImports);
 
 struct ClientSignatureHelpConfiguration
 {
@@ -97,4 +99,4 @@ struct ClientConfiguration
     ClientSignatureHelpConfiguration signatureHelp;
 };
 NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE_WITH_DEFAULT(
-    ClientConfiguration, autocompleteEnd, ignoreGlobs, sourcemap, diagnostics, types, inlayHints, hover, completion, signatureHelp);
+    ClientConfiguration, autocompleteEnd, ignoreGlobs, sourcemap, diagnostics, types, inlayHints, hover, completion, signatureHelp);

src/include/LSP/Workspace.hpp

@@ -59,6 +59,8 @@ class WorkspaceFolder
 
 private:
     void endAutocompletion(const lsp::CompletionParams& params);
+    void suggestImports(const Luau::ModuleName& moduleName, const Luau::Position& position, const ClientConfiguration& config,
+        std::vector<lsp::CompletionItem>& result);
 
 public:
     std::vector<lsp::CompletionItem> completion(const lsp::CompletionParams& params);

src/operations/Completion.cpp

@@ -113,6 +113,156 @@ void WorkspaceFolder::endAutocompletion(const lsp::CompletionParams& params)
     }
 }
 
+bool isGetService(const Luau::AstExpr* expr)
+{
+    if (auto call = expr->as<Luau::AstExprCall>())
+        if (auto index = call->func->as<Luau::AstExprIndexName>())
+            if (index->index == "GetService")
+                if (auto name = index->expr->as<Luau::AstExprGlobal>())
+                    if (name->name == "game")
+                        return true;
+
+    return false;
+}
+
+struct ImportLocationVisitor : public Luau::AstVisitor
+{
+    std::unordered_map<std::string, size_t> serviceLineMap;
+
+    bool visit(Luau::AstStatLocal* local) override
+    {
+        if (local->vars.size != 1)
+            return false;
+
+        auto localName = local->vars.data[0];
+        auto expr = local->values.data[0];
+
+        if (!localName || !expr)
+            return false;
+
+        if (isGetService(expr))
+            serviceLineMap.emplace(std::string(localName->name.value), localName->location.begin.line);
+
+        return false;
+    }
+
+    bool visit(Luau::AstStatBlock* block) override
+    {
+        for (Luau::AstStat* stat : block->body)
+        {
+            stat->visit(this);
+        }
+
+        return false;
+    }
+};
+
+/// Attempts to retrieve a list of service names by inspecting the global type definitions
+static std::vector<std::string> getServiceNames(const Luau::ScopePtr scope)
+{
+    std::vector<std::string> services;
+
+    if (auto dataModelType = scope->lookupType("ServiceProvider"))
+    {
+        if (auto ctv = Luau::get<Luau::ClassTypeVar>(dataModelType->type))
+        {
+            if (auto getService = Luau::lookupClassProp(ctv, "GetService"))
+            {
+                if (auto itv = Luau::get<Luau::IntersectionTypeVar>(getService->type))
+                {
+                    for (auto part : itv->parts)
+                    {
+                        if (auto ftv = Luau::get<Luau::FunctionTypeVar>(part))
+                        {
+                            auto it = Luau::begin(ftv->argTypes);
+                            auto end = Luau::end(ftv->argTypes);
+
+                            if (it != end && ++it != end)
+                            {
+                                if (auto stv = Luau::get<Luau::SingletonTypeVar>(*it))
+                                {
+                                    if (auto ss = Luau::get<Luau::StringSingleton>(stv))
+                                    {
+                                        services.emplace_back(ss->value);
+                                    }
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    return services;
+}
+
+void WorkspaceFolder::suggestImports(
+    const Luau::ModuleName& moduleName, const Luau::Position& position, const ClientConfiguration& config, std::vector<lsp::CompletionItem>& result)
+{
+    auto sourceModule = frontend.getSourceModule(moduleName);
+    auto module = frontend.moduleResolverForAutocomplete.getModule(moduleName);
+    if (!sourceModule || !module)
+        return;
+
+    // If in roblox mode - suggest services
+    if (config.types.roblox)
+    {
+        auto scope = Luau::findScopeAtPosition(*module, position);
+        if (!scope)
+            return;
+
+        // Place after any hot comments and TODO: already imported services
+        size_t minimumLineNumber = 0;
+        for (auto hotComment : sourceModule->hotcomments)
+        {
+            if (!hotComment.header)
+                continue;
+            if (hotComment.location.begin.line >= minimumLineNumber)
+                minimumLineNumber = hotComment.location.begin.line + 1;
+        }
+
+        ImportLocationVisitor visitor;
+        visitor.visit(sourceModule->root);
+
+        auto services = getServiceNames(frontend.typeCheckerForAutocomplete.globalScope);
+        for (auto& service : services)
+        {
+            // ASSUMPTION: if the service was defined, it was defined with the exact same name
+            bool isAlreadyDefined = false;
+            size_t lineNumber = minimumLineNumber;
+            for (auto& [definedService, location] : visitor.serviceLineMap)
+            {
+                if (definedService == service)
+                {
+                    isAlreadyDefined = true;
+                    break;
+                }
+
+                if (definedService < service && location >= lineNumber)
+                    lineNumber = location + 1;
+            }
+
+            if (isAlreadyDefined)
+                continue;
+
+            auto importText = "local " + service + " = game:GetService(\"" + service + "\")\n";
+
+            lsp::CompletionItem item;
+            item.label = service;
+            item.kind = lsp::CompletionItemKind::Class;
+            item.detail = "Auto-import";
+            item.documentation = {lsp::MarkupKind::Markdown, codeBlock("lua", importText)};
+            item.insertText = service;
+
+            lsp::Position placement{lineNumber, 0};
+            item.additionalTextEdits.emplace_back(lsp::TextEdit{{placement, placement}, importText});
+
+            result.emplace_back(item);
+        }
+    }
+}
+
 std::vector<lsp::CompletionItem> WorkspaceFolder::completion(const lsp::CompletionParams& params)
 {
     auto config = client->getConfiguration(rootUri);
@@ -127,7 +277,9 @@ std::vector<lsp::CompletionItem> WorkspaceFolder::completion(const lsp::Completi
         return {};
     }
 
-    auto result = Luau::autocomplete(frontend, fileResolver.getModuleName(params.textDocument.uri), convertPosition(params.position), nullCallback);
+    auto moduleName = fileResolver.getModuleName(params.textDocument.uri);
+    auto position = convertPosition(params.position);
+    auto result = Luau::autocomplete(frontend, moduleName, position, nullCallback);
     std::vector<lsp::CompletionItem> items;
 
     for (auto& [name, entry] : result.entryMap)
@@ -255,6 +407,12 @@ std::vector<lsp::CompletionItem> WorkspaceFolder::completion(const lsp::Completi
         items.emplace_back(item);
     }
 
+    if (config.completion.suggestImports &&
+        (result.context == Luau::AutocompleteContext::Expression || result.context == Luau::AutocompleteContext::Statement))
+    {
+        suggestImports(moduleName, position, config, items);
+    }
+
     return items;
 }