WIP: codegen instructions

CHANGELOG.md

@@ -9,6 +9,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 ### Added
 
 - The VSCode extension now registers a JSON schema for `.luaurc` files, providing simple diagnostics and intellisense ([#850](https://github.com/JohnnyMorganz/luau-lsp/pull/850))
+- Added command `Luau: Compute CodeGen instructions for file` to emit annotated codegen instructions, similar to the bytecode command. External editors can implement this by using the `luau-lsp/codeGen` request. ([#617](https://github.com/JohnnyMorganz/luau-lsp/issues/617))
 
 ### Changed
 

CMakeLists.txt

@@ -138,7 +138,7 @@ set(EXTERN_INCLUDES extern/json/include extern/glob/single_include extern/argpar
 target_compile_features(Luau.LanguageServer PUBLIC cxx_std_17)
 target_compile_options(Luau.LanguageServer PRIVATE ${LUAU_LSP_OPTIONS})
 target_include_directories(Luau.LanguageServer PUBLIC src/include ${EXTERN_INCLUDES})
-target_link_libraries(Luau.LanguageServer PRIVATE Luau.Ast Luau.Analysis Luau.Compiler)
+target_link_libraries(Luau.LanguageServer PRIVATE Luau.Ast Luau.Analysis Luau.Compiler Luau.VM)
 
 set_target_properties(Luau.LanguageServer.CLI PROPERTIES OUTPUT_NAME luau-lsp)
 target_compile_features(Luau.LanguageServer.CLI PUBLIC cxx_std_17)

editors/README.md

@@ -158,11 +158,12 @@ Further Reference:
 
 ## Optional: Bytecode generation
 
-The Language server implements support for computing file-level textual bytecode and source code remarks, for lower level debugging features.
+The Language server implements support for computing file-level textual bytecode, source code remarks, and codegen instructions for lower level debugging features.
 
 A custom LSP request message is implemented:
 
 - `luau-lsp/bytecode`: `{ textDocument: TextDocumentIdentifier, optimizationLevel: number }`, returns `string` - textual bytecode output
 - `luau-lsp/compilerRemarks`: `{ textDocument: TextDocumentIdentifier, optimizationLevel: number }`, returns `string` - source code with inline remarks as comments
+- `luau-lsp/codeGen`: `{ textDocument: TextDocumentIdentifier, optimizationLevel: number }`, returns `string` - annotated codegen instructions
 
 You can implement this request via a custom command to surface this information in your editor

editors/code/package.json

@@ -120,6 +120,11 @@
         "title": "Luau: Compute Compiler Remarks for file",
         "enablement": "resourceLangId == luau"
       },
+      {
+        "command": "luau-lsp.computeCodeGen",
+        "title": "Luau: Compute CodeGen instructions for file",
+        "enablement": "resourceLangId == luau"
+      },
       {
         "command": "luau-lsp.flushTimeTrace",
         "title": "Luau: Flush Time Trace Events to tracing.json"

editors/code/src/bytecode.ts

@@ -8,6 +8,7 @@ import {
 
 export const BYTECODE_SCHEME = "luau-bytecode";
 export const COMPILER_REMARKS_SCHEME = "luau-remarks";
+export const CODEGEN_SCHEME = "luau-codegen";
 
 enum OptimizationLevel {
   None = 0,
@@ -35,6 +36,17 @@ export const ComputeCompilerRemarksRequest = new RequestType<
   void
 >("luau-lsp/compilerRemarks");
 
+export type CodeGenParams = {
+  textDocument: TextDocumentIdentifier;
+  optimizationLevel: OptimizationLevel;
+};
+
+export const ComputeCodeGenRequest = new RequestType<
+  CodeGenParams,
+  string,
+  void
+>("luau-lsp/codeGen");
+
 export const getOptimizationLevel = async (): Promise<OptimizationLevel> => {
   const optimizationLevel = await vscode.window.showQuickPick(
     [
@@ -201,3 +213,17 @@ export const registerComputeCompilerRemarks = (
     ComputeCompilerRemarksRequest,
   );
 };
+
+export const registerComputeCodeGen = (
+  context: vscode.ExtensionContext,
+  client: LanguageClient,
+): vscode.Disposable[] => {
+  return getBytecodeInfo(
+    context,
+    client,
+    "luau-lsp.computeCodeGen",
+    CODEGEN_SCHEME,
+    "codeGen",
+    ComputeCodeGenRequest,
+  );
+};

editors/code/src/extension.ts

@@ -10,6 +10,7 @@ import {
 
 import {
   registerComputeBytecode,
+  registerComputeCodeGen,
   registerComputeCompilerRemarks,
 } from "./bytecode";
 
@@ -239,6 +240,7 @@ const startLanguageServer = async (context: vscode.ExtensionContext) => {
 
   clientDisposables.push(...registerComputeBytecode(context, client));
   clientDisposables.push(...registerComputeCompilerRemarks(context, client));
+  clientDisposables.push(...registerComputeCodeGen(context, client));
 
   console.log("LSP Setup");
   await client.start();

src/LanguageServer.cpp

@@ -258,14 +258,20 @@ void LanguageServer::onRequest(const id_type& id, const std::string& method, std
         auto workspace = findWorkspace(params.textDocument.uri);
         response = workspace->bytecode(params);
     }
-
     else if (method == "luau-lsp/compilerRemarks")
     {
         ASSERT_PARAMS(baseParams, "luau-lsp/compilerRemarks")
         auto params = baseParams->get<lsp::CompilerRemarksParams>();
         auto workspace = findWorkspace(params.textDocument.uri);
         response = workspace->compilerRemarks(params);
     }
+    else if (method == "luau-lsp/codegen")
+    {
+        ASSERT_PARAMS(baseParams, "luau-lsp/codeGen")
+        auto params = baseParams->get<lsp::CodegenParams>();
+        auto workspace = findWorkspace(params.textDocument.uri);
+        response = workspace->codeGen(params);
+    }
     else
     {
         throw JsonRpcException(lsp::ErrorCode::MethodNotFound, "method not found / supported: " + method);

src/include/LSP/Workspace.hpp

@@ -130,6 +130,7 @@ class WorkspaceFolder
 
     lsp::BytecodeResult bytecode(const lsp::BytecodeParams& params);
     lsp::CompilerRemarksResult compilerRemarks(const lsp::CompilerRemarksParams& params);
+    lsp::CodegenResult codeGen(const lsp::CodegenParams& params);
 
     bool isNullWorkspace() const
     {

src/include/Protocol/Extensions.hpp

@@ -28,4 +28,13 @@ struct CompilerRemarksParams
 NLOHMANN_DEFINE_OPTIONAL(CompilerRemarksParams, textDocument, optimizationLevel)
 
 using CompilerRemarksResult = std::string;
+
+struct CodegenParams
+{
+    TextDocumentIdentifier textDocument;
+    CompilerRemarksOptimizationLevel optimizationLevel = CompilerRemarksOptimizationLevel::O1;
+};
+NLOHMANN_DEFINE_OPTIONAL(CodegenParams, textDocument, optimizationLevel)
+
+using CodegenResult = std::string;
 } // namespace lsp

src/operations/Bytecode.cpp

@@ -1,17 +1,44 @@
+#include "../../luau/CodeGen/include/Luau/CodeGen.h"
 #include "LSP/Workspace.hpp"
 #include "Luau/BytecodeBuilder.h"
 #include "Luau/Parser.h"
 #include "Luau/Compiler.h"
 
+#include "lua.h"
+#include "lualib.h"
+
 static std::string constructError(const std::string& type, const Luau::Location& location, const std::string& message)
 {
     return type + "(" + std::to_string(location.begin.line + 1) + "," + std::to_string(location.begin.column + 1) + "): " + message + "\n";
 }
 
+static std::string getCodegenAssembly(
+    const char* name,
+    const std::string& bytecode,
+    Luau::CodeGen::AssemblyOptions options
+)
+{
+    std::unique_ptr<lua_State, void (*)(lua_State*)> globalState(luaL_newstate(), lua_close);
+    lua_State* L = globalState.get();
+
+    if (luau_load(L, name, bytecode.data(), bytecode.size(), 0) == 0)
+        return Luau::CodeGen::getAssembly(L, -1, options, nullptr);
+
+    return "Error loading bytecode";
+}
+
+static void annotateInstruction(void* context, std::string& text, int fid, int instpos)
+{
+    Luau::BytecodeBuilder& bcb = *(Luau::BytecodeBuilder*)context;
+
+    bcb.annotateInstruction(text, fid, instpos);
+}
+
 enum class BytecodeOutputType
 {
     Textual,
-    CompilerRemarks
+    CompilerRemarks,
+    CodeGen,
 };
 
 static uint32_t flagsForType(BytecodeOutputType type)
@@ -23,11 +50,14 @@ static uint32_t flagsForType(BytecodeOutputType type)
                Luau::BytecodeBuilder::Dump_Remarks | Luau::BytecodeBuilder::Dump_Types;
     case BytecodeOutputType::CompilerRemarks:
         return Luau::BytecodeBuilder::Dump_Source | Luau::BytecodeBuilder::Dump_Remarks;
+    case BytecodeOutputType::CodeGen:
+        return Luau::BytecodeBuilder::Dump_Code | Luau::BytecodeBuilder::Dump_Source | Luau::BytecodeBuilder::Dump_Locals |
+            Luau::BytecodeBuilder::Dump_Remarks;
     }
     return 0;
 }
 
-static std::string computeBytecodeOutput(const std::string& source, const ClientConfiguration& config, int optimizationLevel, BytecodeOutputType type)
+static std::string computeBytecodeOutput(const Luau::ModuleName& moduleName, const std::string& source, const ClientConfiguration& config, int optimizationLevel, BytecodeOutputType type)
 {
     try
     {
@@ -52,7 +82,20 @@ static std::string computeBytecodeOutput(const std::string& source, const Client
 
         Luau::compileOrThrow(bcb, result, names, options);
 
-        if (type == BytecodeOutputType::Textual)
+        if (type == BytecodeOutputType::CodeGen)
+        {
+            Luau::CodeGen::AssemblyOptions assemblyOptions;
+            // TODO: assemblyOptions target
+            assemblyOptions.outputBinary = false;
+            assemblyOptions.includeAssembly = true;
+            assemblyOptions.includeIr = true;
+            assemblyOptions.includeIrTypes = true;
+            assemblyOptions.includeOutlinedCode = true;
+            assemblyOptions.annotator = annotateInstruction;
+            assemblyOptions.annotatorContext = &bcb;
+            return getCodegenAssembly(moduleName.c_str(), bcb.getBytecode(), assemblyOptions);
+        }
+        else if (type == BytecodeOutputType::Textual)
             return bcb.dumpEverything();
         else
             return bcb.dumpSourceRemarks();
@@ -78,7 +121,7 @@ lsp::CompilerRemarksResult WorkspaceFolder::bytecode(const lsp::BytecodeParams&
         throw JsonRpcException(lsp::ErrorCode::RequestFailed, "No managed text document for " + params.textDocument.uri.toString());
 
     auto config = client->getConfiguration(rootUri);
-    return computeBytecodeOutput(textDocument->getText(), config, params.optimizationLevel, BytecodeOutputType::Textual);
+    return computeBytecodeOutput(moduleName, textDocument->getText(), config, params.optimizationLevel, BytecodeOutputType::Textual);
 }
 
 lsp::CompilerRemarksResult WorkspaceFolder::compilerRemarks(const lsp::CompilerRemarksParams& params)
@@ -89,5 +132,16 @@ lsp::CompilerRemarksResult WorkspaceFolder::compilerRemarks(const lsp::CompilerR
         throw JsonRpcException(lsp::ErrorCode::RequestFailed, "No managed text document for " + params.textDocument.uri.toString());
 
     auto config = client->getConfiguration(rootUri);
-    return computeBytecodeOutput(textDocument->getText(), config, params.optimizationLevel, BytecodeOutputType::CompilerRemarks);
+    return computeBytecodeOutput(moduleName, textDocument->getText(), config, params.optimizationLevel, BytecodeOutputType::CompilerRemarks);
+}
+
+lsp::CompilerRemarksResult WorkspaceFolder::codeGen(const lsp::CodegenParams& params)
+{
+    auto moduleName = fileResolver.getModuleName(params.textDocument.uri);
+    auto textDocument = fileResolver.getTextDocument(params.textDocument.uri);
+    if (!textDocument)
+        throw JsonRpcException(lsp::ErrorCode::RequestFailed, "No managed text document for " + params.textDocument.uri.toString());
+
+    auto config = client->getConfiguration(rootUri);
+    return computeBytecodeOutput(moduleName, textDocument->getText(), config, params.optimizationLevel, BytecodeOutputType::CodeGen);
 }