nixd: add basic goto-def support via libnixf (no eval) (#400)

nixd/include/nixd/Controller/Controller.h

@@ -18,7 +18,31 @@ class Controller : public lspserver::LSPServer {
       PublishDiagnostic;
 
   std::mutex TUsLock;
-  llvm::StringMap<NixTU> TUs;
+  llvm::StringMap<std::shared_ptr<NixTU>> TUs;
+
+  template <class T>
+  std::shared_ptr<NixTU> getTU(std::string File, lspserver::Callback<T> &Reply,
+                               bool Ignore = false) {
+    using lspserver::error;
+    std::lock_guard G(TUsLock);
+    if (!TUs.count(File)) [[unlikely]] {
+      if (!Ignore)
+        Reply(error("cannot find corresponding AST on file {0}", File));
+      return nullptr;
+    }
+    return TUs[File];
+  }
+
+  template <class T>
+  std::shared_ptr<nixf::Node> getAST(const NixTU &TU,
+                                     lspserver::Callback<T> &Reply) {
+    using lspserver::error;
+    if (!TU.ast()) {
+      Reply(error("AST is null on this unit"));
+      return nullptr;
+    }
+    return TU.ast();
+  }
 
   boost::asio::thread_pool Pool;
 
@@ -52,6 +76,9 @@ class Controller : public lspserver::LSPServer {
   void onHover(const lspserver::TextDocumentPositionParams &Params,
                lspserver::Callback<std::optional<lspserver::Hover>> Reply);
 
+  void onDefinition(const lspserver::TextDocumentPositionParams &Params,
+                    lspserver::Callback<lspserver::Location> Reply);
+
   void publishDiagnostics(lspserver::PathRef File,
                           std::optional<int64_t> Version,
                           const std::vector<nixf::Diagnostic> &Diagnostics);

nixd/include/nixd/Controller/NixTU.h

@@ -4,6 +4,8 @@
 
 #include "nixf/Basic/Diagnostic.h"
 #include "nixf/Basic/Nodes/Basic.h"
+#include "nixf/Sema/ParentMap.h"
+#include "nixf/Sema/VariableLookup.h"
 
 #include <memory>
 #include <optional>
@@ -18,20 +20,29 @@ class NixTU {
   std::vector<nixf::Diagnostic> Diagnostics;
   std::shared_ptr<nixf::Node> AST;
   std::optional<util::OwnedRegion> ASTByteCode;
+  std::unique_ptr<nixf::VariableLookupAnalysis> VLA;
+  std::unique_ptr<nixf::ParentMapAnalysis> PMA;
 
 public:
   NixTU() = default;
   NixTU(std::vector<nixf::Diagnostic> Diagnostics,
         std::shared_ptr<nixf::Node> AST,
-        std::optional<util::OwnedRegion> ASTByteCode)
-      : Diagnostics(std::move(Diagnostics)), AST(std::move(AST)),
-        ASTByteCode(std::move(ASTByteCode)) {}
+        std::optional<util::OwnedRegion> ASTByteCode,
+        std::unique_ptr<nixf::VariableLookupAnalysis> VLA);
 
   [[nodiscard]] const std::vector<nixf::Diagnostic> &diagnostics() const {
     return Diagnostics;
   }
 
   [[nodiscard]] const std::shared_ptr<nixf::Node> &ast() const { return AST; }
+
+  [[nodiscard]] const nixf::ParentMapAnalysis *parentMap() const {
+    return PMA.get();
+  }
+
+  [[nodiscard]] const nixf::VariableLookupAnalysis *variableLookup() const {
+    return VLA.get();
+  }
 };
 
 } // namespace nixd

nixd/lib/Controller/CodeAction.cpp

@@ -19,41 +19,39 @@ void Controller::onCodeAction(const lspserver::CodeActionParams &Params,
   std::string File(Params.textDocument.uri.file());
   Range Range = Params.range;
   auto Action = [Reply = std::move(Reply), File, Range, this]() mutable {
-    std::vector<nixf::Diagnostic> Diagnostics;
-    {
-      std::lock_guard TU(TUsLock);
-      Diagnostics = TUs[File].diagnostics();
-    }
-    std::vector<CodeAction> Actions;
-    Actions.reserve(Diagnostics.size());
-    for (const nixf::Diagnostic &D : Diagnostics) {
-      auto DRange = toLSPRange(D.range());
-      if (!Range.overlap(DRange))
-        continue;
-
-      // Add fixes.
-      for (const nixf::Fix &F : D.fixes()) {
-        std::vector<TextEdit> Edits;
-        Edits.reserve(F.edits().size());
-        for (const nixf::TextEdit &TE : F.edits()) {
-          Edits.emplace_back(TextEdit{
-              .range = toLSPRange(TE.oldRange()),
-              .newText = std::string(TE.newText()),
+    if (auto TU = getTU(File, Reply, /*Ignore=*/true)) {
+      std::vector<nixf::Diagnostic> Diagnostics = TU->diagnostics();
+      std::vector<CodeAction> Actions;
+      Actions.reserve(Diagnostics.size());
+      for (const nixf::Diagnostic &D : Diagnostics) {
+        auto DRange = toLSPRange(D.range());
+        if (!Range.overlap(DRange))
+          continue;
+
+        // Add fixes.
+        for (const nixf::Fix &F : D.fixes()) {
+          std::vector<TextEdit> Edits;
+          Edits.reserve(F.edits().size());
+          for (const nixf::TextEdit &TE : F.edits()) {
+            Edits.emplace_back(TextEdit{
+                .range = toLSPRange(TE.oldRange()),
+                .newText = std::string(TE.newText()),
+            });
+          }
+          using Changes = std::map<std::string, std::vector<TextEdit>>;
+          std::string FileURI = URIForFile::canonicalize(File, File).uri();
+          WorkspaceEdit WE{.changes = Changes{
+                               {std::move(FileURI), std::move(Edits)},
+                           }};
+          Actions.emplace_back(CodeAction{
+              .title = F.message(),
+              .kind = std::string(CodeAction::QUICKFIX_KIND),
+              .edit = std::move(WE),
           });
         }
-        using Changes = std::map<std::string, std::vector<TextEdit>>;
-        std::string FileURI = URIForFile::canonicalize(File, File).uri();
-        WorkspaceEdit WE{.changes = Changes{
-                             {std::move(FileURI), std::move(Edits)},
-                         }};
-        Actions.emplace_back(CodeAction{
-            .title = F.message(),
-            .kind = std::string(CodeAction::QUICKFIX_KIND),
-            .edit = std::move(WE),
-        });
       }
+      Reply(std::move(Actions));
     }
-    Reply(std::move(Actions));
   };
   boost::asio::post(Pool, std::move(Action));
 }

nixd/lib/Controller/Convert.cpp

@@ -21,6 +21,10 @@ lspserver::Position toLSPPosition(const nixf::LexerCursor &P) {
                              static_cast<int>(P.column())};
 }
 
+nixf::Position toNixfPosition(const lspserver::Position &P) {
+  return {P.line, P.character};
+}
+
 lspserver::Range toLSPRange(const nixf::LexerCursorRange &R) {
   return lspserver::Range{toLSPPosition(R.lCur()), toLSPPosition(R.rCur())};
 }

nixd/lib/Controller/Convert.h

@@ -12,6 +12,8 @@ namespace nixd {
 
 lspserver::Position toLSPPosition(const nixf::LexerCursor &P);
 
+nixf::Position toNixfPosition(const lspserver::Position &P);
+
 lspserver::Range toLSPRange(const nixf::LexerCursorRange &R);
 
 int getLSPSeverity(nixf::Diagnostic::DiagnosticKind Kind);

nixd/lib/Controller/Definition.cpp

@@ -0,0 +1,80 @@
+/// \file
+/// \brief Implementation of [Go to Definition]
+/// [Go to Definition]:
+/// https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_definition
+
+#include "Convert.h"
+
+#include "nixd/Controller/Controller.h"
+
+#include <llvm/Support/Error.h>
+
+#include <boost/asio/post.hpp>
+
+using namespace nixd;
+using namespace lspserver;
+
+namespace {
+
+void gotoDefinition(const NixTU &TU, const nixf::Node &AST, nixf::Position Pos,
+                    URIForFile URI, Callback<Location> &Reply) {
+  using LookupResult = nixf::VariableLookupAnalysis::LookupResult;
+  using ResultKind = nixf::VariableLookupAnalysis::LookupResultKind;
+
+  const nixf::Node *N = AST.descend({Pos, Pos});
+  if (!N) [[unlikely]] {
+    Reply(error("cannot find AST node on given position"));
+    return;
+  }
+
+  const nixf::ParentMapAnalysis *PMA = TU.parentMap();
+  assert(PMA && "ParentMap should not be null as AST is not null");
+
+  const nixf::Node *Var = PMA->upTo(*N, nixf::Node::NK_ExprVar);
+  if (!Var) [[unlikely]] {
+    Reply(error("cannot find variable on given position"));
+    return;
+  }
+  assert(Var->kind() == nixf::Node::NK_ExprVar);
+
+  // OK, this is an variable. Lookup it in VLA entries.
+  const nixf::VariableLookupAnalysis *VLA = TU.variableLookup();
+  if (!VLA) [[unlikely]] {
+    Reply(error("cannot get variable analysis for nix unit"));
+    return;
+  }
+
+  LookupResult Result = VLA->query(static_cast<const nixf::ExprVar &>(*Var));
+  if (Result.Kind == ResultKind::Undefined) {
+    Reply(error("this varaible is undefined"));
+    return;
+  }
+
+  if (Result.Def->isBuiltin()) {
+    Reply(error("this is a builtin variable"));
+    return;
+  }
+
+  assert(Result.Def->syntax());
+
+  Reply(Location{
+      .uri = std::move(URI),
+      .range = toLSPRange(Result.Def->syntax()->range()),
+  });
+}
+
+} // namespace
+
+void Controller::onDefinition(const TextDocumentPositionParams &Params,
+                              Callback<Location> Reply) {
+  auto Action = [Reply = std::move(Reply), URI = Params.textDocument.uri,
+                 Pos = toNixfPosition(Params.position), this]() mutable {
+    std::string File(URI.file());
+    if (std::shared_ptr<NixTU> TU = getTU(File, Reply)) [[likely]] {
+      if (std::shared_ptr<nixf::Node> AST = getAST(*TU, Reply)) [[likely]] {
+        gotoDefinition(*TU, *AST, Pos, std::move(URI), Reply);
+      }
+    }
+  };
+  boost::asio::post(Pool, std::move(Action));
+}

nixd/lib/Controller/Hover.cpp

@@ -24,23 +24,25 @@ void Controller::onHover(const TextDocumentPositionParams &Params,
   auto Action = [Reply = std::move(Reply),
                  File = std::string(Params.textDocument.uri.file()),
                  RawPos = Params.position, this]() mutable {
-    std::lock_guard G(TUsLock);
-    const std::shared_ptr<nixf::Node> &AST = TUs[File].ast();
-    nixf::Position Pos{RawPos.line, RawPos.character};
-    const nixf::Node *N = AST->descend({Pos, Pos});
-    if (!N) {
-      Reply(std::nullopt);
-      return;
+    if (std::shared_ptr<NixTU> TU = getTU(File, Reply)) [[likely]] {
+      if (std::shared_ptr<nixf::Node> AST = getAST(*TU, Reply)) [[likely]] {
+        nixf::Position Pos{RawPos.line, RawPos.character};
+        const nixf::Node *N = AST->descend({Pos, Pos});
+        if (!N) {
+          Reply(std::nullopt);
+          return;
+        }
+        std::string Name = N->name();
+        Reply(Hover{
+            .contents =
+                MarkupContent{
+                    .kind = MarkupKind::Markdown,
+                    .value = "`" + Name + "`",
+                },
+            .range = toLSPRange(N->range()),
+        });
+      }
     }
-    std::string Name = N->name();
-    Reply(Hover{
-        .contents =
-            MarkupContent{
-                .kind = MarkupKind::Markdown,
-                .value = "`" + Name + "`",
-            },
-        .range = toLSPRange(N->range()),
-    });
   };
   boost::asio::post(Pool, std::move(Action));
 }

nixd/lib/Controller/LifeTime.cpp

@@ -36,6 +36,7 @@ void Controller::
                {"resolveProvider", false},
            },
        },
+       {"definitionProvider", true},
        {"hoverProvider", true}},
   };
 

nixd/lib/Controller/NixTU.cpp

@@ -0,0 +1,15 @@
+#include "nixd/Controller/NixTU.h"
+
+using namespace nixd;
+
+NixTU::NixTU(std::vector<nixf::Diagnostic> Diagnostics,
+             std::shared_ptr<nixf::Node> AST,
+             std::optional<util::OwnedRegion> ASTByteCode,
+             std::unique_ptr<nixf::VariableLookupAnalysis> VLA)
+    : Diagnostics(std::move(Diagnostics)), AST(std::move(AST)),
+      ASTByteCode(std::move(ASTByteCode)), VLA(std::move(VLA)) {
+  if (this->AST) {
+    PMA = std::make_unique<nixf::ParentMapAnalysis>();
+    PMA->runOnAST(*this->AST);
+  }
+}

nixd/lib/Controller/Support.cpp

@@ -42,12 +42,15 @@ void Controller::actOnDocumentAdd(PathRef File,
     if (!AST) {
       std::lock_guard G(TUsLock);
       publishDiagnostics(File, Version, Diagnostics);
-      TUs[File] = NixTU(std::move(Diagnostics), std::move(AST), std::nullopt);
+      TUs.insert_or_assign(File,
+                           std::make_shared<NixTU>(std::move(Diagnostics),
+                                                   std::move(AST), std::nullopt,
+                                                   /*VLA=*/nullptr));
       return;
     }
 
-    nixf::VariableLookupAnalysis VLA(Diagnostics);
-    VLA.runOnAST(*AST);
+    auto VLA = std::make_unique<nixf::VariableLookupAnalysis>(Diagnostics);
+    VLA->runOnAST(*AST);
 
     publishDiagnostics(File, Version, Diagnostics);
 
@@ -58,7 +61,9 @@ void Controller::actOnDocumentAdd(PathRef File,
     if (Buf.empty()) {
       lspserver::log("empty AST for {0}", File);
       std::lock_guard G(TUsLock);
-      TUs[File] = NixTU(std::move(Diagnostics), std::move(AST), std::nullopt);
+      TUs.insert_or_assign(
+          File, std::make_shared<NixTU>(std::move(Diagnostics), std::move(AST),
+                                        std::nullopt, std::move(VLA)));
       return;
     }
 
@@ -77,8 +82,11 @@ void Controller::actOnDocumentAdd(PathRef File,
                    Buf.size());
 
     std::lock_guard G(TUsLock);
-    TUs[File] = NixTU(std::move(Diagnostics), std::move(AST),
-                      util::OwnedRegion{std::move(Shm), std::move(Region)});
+    TUs.insert_or_assign(
+        File, std::make_shared<NixTU>(
+                  std::move(Diagnostics), std::move(AST),
+                  util::OwnedRegion{std::move(Shm), std::move(Region)},
+                  std::move(VLA)));
 
     if (Eval) {
       Eval->RegisterBC(rpc::RegisterBCParams{
@@ -91,8 +99,8 @@ void Controller::actOnDocumentAdd(PathRef File,
     // So just invoke the action here.
     Action();
   } else {
-    // Otherwise we may want to concurently parse & serialize the file, so post
-    // it to the thread pool.
+    // Otherwise we may want to concurently parse & serialize the file, so
+    // post it to the thread pool.
     boost::asio::post(Pool, std::move(Action));
   }
 }
@@ -115,6 +123,8 @@ Controller::Controller(std::unique_ptr<lspserver::InboundPort> In,
                            &Controller::onDocumentDidClose);
 
   // Language Features
+  Registry.addMethod("textDocument/definition", this,
+                     &Controller::onDefinition);
   Registry.addMethod("textDocument/codeAction", this,
                      &Controller::onCodeAction);
   Registry.addMethod("textDocument/hover", this, &Controller::onHover);

nixd/meson.build

@@ -6,10 +6,12 @@ libnixd_lib = library(
     'nixd',
     'lib/Controller/CodeAction.cpp',
     'lib/Controller/Convert.cpp',
+    'lib/Controller/Definition.cpp',
     'lib/Controller/Diagnostics.cpp',
     'lib/Controller/EvalClient.cpp',
     'lib/Controller/Hover.cpp',
     'lib/Controller/LifeTime.cpp',
+    'lib/Controller/NixTU.cpp',
     'lib/Controller/Support.cpp',
     'lib/Controller/TextDocumentSync.cpp',
     'lib/Eval/EvalProvider.cpp',