nixd: support find references

nixd/include/nixd/Controller/Controller.h

@@ -79,6 +79,10 @@ class Controller : public lspserver::LSPServer {
   void onDefinition(const lspserver::TextDocumentPositionParams &Params,
                     lspserver::Callback<lspserver::Location> Reply);
 
+  void
+  onReferences(const lspserver::TextDocumentPositionParams &Params,
+               lspserver::Callback<std::vector<lspserver::Location>> Reply);
+
   void publishDiagnostics(lspserver::PathRef File,
                           std::optional<int64_t> Version,
                           const std::vector<nixf::Diagnostic> &Diagnostics);

nixd/lib/Controller/Definition.cpp

@@ -3,75 +3,102 @@
 /// [Go to Definition]:
 /// https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_definition
 
+#include "Definition.h"
 #include "Convert.h"
 
 #include "nixd/Controller/Controller.h"
+#include "nixf/Sema/VariableLookup.h"
 
 #include <llvm/Support/Error.h>
 
 #include <boost/asio/post.hpp>
 
 using namespace nixd;
+using namespace nixf;
 using namespace lspserver;
+using namespace llvm;
+
+using LookupResult = VariableLookupAnalysis::LookupResult;
+using ResultKind = VariableLookupAnalysis::LookupResultKind;
 
 namespace {
 
-void gotoDefinition(const NixTU &TU, const nixf::Node &AST, nixf::Position Pos,
+void gotoDefinition(const NixTU &TU, const 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});
+  const 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();
+  const 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();
+  const 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 (Expected<const Definition &> ExpDef = findDefinition(*N, *PMA, *VLA)) {
+    assert(ExpDef->syntax());
+    Reply(Location{
+        .uri = std::move(URI),
+        .range = toLSPRange(ExpDef->syntax()->range()),
+    });
+  } else {
+    Reply(ExpDef.takeError());
   }
+}
 
-  if (Result.Def->isBuiltin()) {
-    Reply(error("this is a builtin variable"));
-    return;
-  }
+const Definition *findSelfDefinition(const Node &N,
+                                     const ParentMapAnalysis &PMA,
+                                     const VariableLookupAnalysis &VLA) {
+  // If "N" is a definition itself, just return it.
+  if (const Definition *Def = VLA.toDef(N))
+    return Def;
 
-  assert(Result.Def->syntax());
+  // If N is inside an attrset, it maybe an "AttrName", let's look for it.
+  const Node *Parent = PMA.query(N);
+  if (Parent && Parent->kind() == Node::NK_AttrName)
+    return VLA.toDef(*Parent);
 
-  Reply(Location{
-      .uri = std::move(URI),
-      .range = toLSPRange(Result.Def->syntax()->range()),
-  });
+  return nullptr;
 }
 
 } // namespace
 
+Expected<const Definition &>
+nixd::findDefinition(const Node &N, const ParentMapAnalysis &PMA,
+                     const VariableLookupAnalysis &VLA) {
+  const Node *Var = PMA.upTo(N, Node::NK_ExprVar);
+  if (!Var) [[unlikely]] {
+    if (const Definition *Def = findSelfDefinition(N, PMA, VLA))
+      return *Def;
+    return error("cannot find variable on given position");
+  }
+  assert(Var->kind() == Node::NK_ExprVar);
+  LookupResult Result = VLA.query(static_cast<const ExprVar &>(*Var));
+
+  assert(Result.Def);
+
+  if (Result.Kind == ResultKind::Undefined)
+    return error("this varaible is undefined");
+
+  if (Result.Def->isBuiltin())
+    return error("this is a builtin variable");
+
+  return *Result.Def;
+}
+
 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]] {
+      if (std::shared_ptr<Node> AST = getAST(*TU, Reply)) [[likely]] {
         gotoDefinition(*TU, *AST, Pos, std::move(URI), Reply);
       }
     }

nixd/lib/Controller/Definition.h

@@ -0,0 +1,15 @@
+#pragma once
+
+#include "nixf/Sema/ParentMap.h"
+#include "nixf/Sema/VariableLookup.h"
+
+#include <llvm/Support/Error.h>
+
+namespace nixd {
+
+/// \brief Heuristically find definition on some node
+llvm::Expected<const nixf::Definition &>
+findDefinition(const nixf::Node &N, const nixf::ParentMapAnalysis &PMA,
+               const nixf::VariableLookupAnalysis &VLA);
+
+} // namespace nixd

nixd/lib/Controller/FindReferences.cpp

@@ -0,0 +1,71 @@
+/// \file
+/// \brief This implements [Find References].
+/// [Find References]:
+/// https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_references
+
+#include "Convert.h"
+#include "Definition.h"
+
+#include "nixd/Controller/Controller.h"
+
+#include <boost/asio/post.hpp>
+#include <lspserver/Protocol.h>
+#include <nixf/Sema/ParentMap.h>
+#include <nixf/Sema/VariableLookup.h>
+
+using namespace lspserver;
+using namespace nixd;
+using namespace llvm;
+using namespace nixf;
+
+namespace {
+
+Error findReferences(const nixf::Node &Desc, const ParentMapAnalysis &PMA,
+                     const VariableLookupAnalysis &VLA, const URIForFile &URI,
+                     std::vector<Location> &Locations) {
+
+  // Two steps.
+  //   1. Find some "definition" for this node.
+  //   2. Find all "uses", and construct the vector.
+
+  // Find "definition"
+  if (auto Def = findDefinition(Desc, PMA, VLA)) {
+    // OK, iterate all uses.
+    for (const auto *Use : Def->uses()) {
+      assert(Use);
+      Locations.emplace_back(Location{
+          .uri = URI,
+          .range = toLSPRange(Use->range()),
+      });
+    }
+    return Error::success();
+  }
+  return error("Cannot find definition of this node");
+}
+
+} // namespace
+
+void Controller::onReferences(const TextDocumentPositionParams &Params,
+                              Callback<std::vector<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]] {
+        const nixf::Node *Desc = AST->descend({Pos, Pos});
+        if (!Desc) {
+          Reply(error("cannot find corresponding node on given position"));
+          return;
+        }
+        std::vector<Location> Locations;
+        if (auto Err = findReferences(*Desc, *TU->parentMap(),
+                                      *TU->variableLookup(), URI, Locations)) {
+          Reply(std::move(Err));
+          return;
+        }
+        Reply(std::move(Locations));
+      }
+    }
+  };
+  boost::asio::post(Pool, std::move(Action));
+}

nixd/lib/Controller/LifeTime.cpp

@@ -37,6 +37,7 @@ void Controller::
            },
        },
        {"definitionProvider", true},
+       {"referencesProvider", true},
        {"hoverProvider", true}},
   };
 

nixd/lib/Controller/Support.cpp

@@ -125,6 +125,8 @@ Controller::Controller(std::unique_ptr<lspserver::InboundPort> In,
   // Language Features
   Registry.addMethod("textDocument/definition", this,
                      &Controller::onDefinition);
+  Registry.addMethod("textDocument/references", this,
+                     &Controller::onReferences);
   Registry.addMethod("textDocument/codeAction", this,
                      &Controller::onCodeAction);
   Registry.addMethod("textDocument/hover", this, &Controller::onHover);

nixd/meson.build

@@ -9,6 +9,7 @@ libnixd_lib = library(
     'lib/Controller/Definition.cpp',
     'lib/Controller/Diagnostics.cpp',
     'lib/Controller/EvalClient.cpp',
+    'lib/Controller/FindReferences.cpp',
     'lib/Controller/Hover.cpp',
     'lib/Controller/LifeTime.cpp',
     'lib/Controller/NixTU.cpp',

nixd/tools/nixd/test/initialize.md

@@ -34,6 +34,7 @@ CHECK-NEXT:         "resolveProvider": false
 CHECK-NEXT:       },
 CHECK-NEXT:       "definitionProvider": true,
 CHECK-NEXT:       "hoverProvider": true,
+CHECK-NEXT:       "referencesProvider": true,
 CHECK-NEXT:       "textDocumentSync": {
 CHECK-NEXT:         "change": 2,
 CHECK-NEXT:         "openClose": true,

nixd/tools/nixd/test/references-lambda.md

@@ -0,0 +1,93 @@
+# RUN: nixd --lit-test < %s | FileCheck %s
+
+<-- initialize(0)
+
+```json
+{
+   "jsonrpc":"2.0",
+   "id":0,
+   "method":"initialize",
+   "params":{
+      "processId":123,
+      "rootPath":"",
+      "capabilities":{
+      },
+      "trace":"off"
+   }
+}
+```
+
+
+<-- textDocument/didOpen
+
+```json
+{
+   "jsonrpc":"2.0",
+   "method":"textDocument/didOpen",
+   "params":{
+      "textDocument":{
+         "uri":"file:///basic.nix",
+         "languageId":"nix",
+         "version":1,
+         "text":"arg @ { foo, bar }: arg + foo + bar + arg"
+      }
+   }
+}
+```
+
+<-- textDocument/references(2)
+
+
+```json
+{
+   "jsonrpc":"2.0",
+   "id":2,
+   "method":"textDocument/references",
+   "params":{
+      "textDocument":{
+         "uri":"file:///basic.nix"
+      },
+      "position":{
+        "line": 0,
+        "character":0
+      }
+   }
+}
+```
+
+```
+     CHECK: "id": 2,
+CHECK-NEXT: "jsonrpc": "2.0",
+CHECK-NEXT: "result": [
+CHECK-NEXT:   {
+CHECK-NEXT:     "range": {
+CHECK-NEXT:       "end": {
+CHECK-NEXT:         "character": 23,
+CHECK-NEXT:         "line": 0
+CHECK-NEXT:       },
+CHECK-NEXT:       "start": {
+CHECK-NEXT:         "character": 20,
+CHECK-NEXT:         "line": 0
+CHECK-NEXT:       }
+CHECK-NEXT:     },
+CHECK-NEXT:     "uri": "file:///basic.nix"
+CHECK-NEXT:   },
+CHECK-NEXT:   {
+CHECK-NEXT:     "range": {
+CHECK-NEXT:       "end": {
+CHECK-NEXT:         "character": 41,
+CHECK-NEXT:         "line": 0
+CHECK-NEXT:       },
+CHECK-NEXT:       "start": {
+CHECK-NEXT:         "character": 38,
+CHECK-NEXT:         "line": 0
+CHECK-NEXT:       }
+CHECK-NEXT:     },
+CHECK-NEXT:     "uri": "file:///basic.nix"
+CHECK-NEXT:   }
+CHECK-NEXT: ]
+```
+
+```json
+{"jsonrpc":"2.0","method":"exit"}
+```