nixd: support find references (#402)

libnixf/include/nixf/Sema/VariableLookup.h

@@ -86,12 +86,16 @@ class VariableLookupAnalysis {
     std::shared_ptr<const Definition> Def;
   };
 
+  using ToDefMap = std::map<const Node *, std::shared_ptr<Definition>>;
+
 private:
   std::vector<Diagnostic> &Diags;
 
   std::map<const Node *, std::shared_ptr<Definition>>
       WithDefs; // record with ... ; users.
 
+  ToDefMap ToDef;
+
   void lookupVar(const ExprVar &Var, const std::shared_ptr<EnvNode> &Env);
 
   std::shared_ptr<EnvNode> dfsAttrs(const SemaAttrs &SA,
@@ -118,11 +122,33 @@ class VariableLookupAnalysis {
 public:
   VariableLookupAnalysis(std::vector<Diagnostic> &Diags);
 
+  /// \brief Perform variable lookup analysis (def-use) on AST.
+  /// \note This method should be invoked after any other method called.
+  /// \note The result remains immutable thus it can be shared among threads.
   void runOnAST(const Node &Root);
 
+  /// \brief Query the which name/with binds to specific varaible.
   [[nodiscard]] LookupResult query(const ExprVar &Var) const {
     return Results.at(&Var);
   }
+
+  /// \brief Get definition record for some name.
+  ///
+  /// For some cases, we need to get "definition" record to find all references
+  /// to this definition, on AST.
+  ///
+  /// Thus we need to store AST -> Definition
+  /// There are many pointers on AST, the convention is:
+  ///
+  ///   1. attrname "key" syntax is recorded.
+  //        For static attrs, they are Node::NK_AttrName.
+  ///   2. "with" keyword is recorded.
+  ///   3. Lambda arguments, record its identifier.
+  [[nodiscard]] const Definition *toDef(const Node &N) const {
+    if (ToDef.contains(&N))
+      return ToDef.at(&N).get();
+    return nullptr;
+  }
 };
 
 } // namespace nixf

libnixf/src/Sema/VariableLookup.cpp

@@ -13,10 +13,17 @@ class DefBuilder {
   EnvNode::DefMap Def;
 
 public:
-  void addBuiltin(std::string Name) { add(std::move(Name), nullptr); }
-  void add(std::string Name, const Node *Entry) {
+  void addBuiltin(std::string Name) {
+    // Don't need to record def map for builtins.
+    auto _ = add(std::move(Name), nullptr);
+  }
+
+  [[nodiscard("Record ToDef Map!")]] std::shared_ptr<Definition>
+  add(std::string Name, const Node *Entry) {
     assert(!Def.contains(Name));
-    Def.insert({std::move(Name), std::make_shared<Definition>(Entry)});
+    auto NewDef = std::make_shared<Definition>(Entry);
+    Def.insert({std::move(Name), NewDef});
+    return NewDef;
   }
 
   EnvNode::DefMap finish() { return std::move(Def); }
@@ -100,14 +107,15 @@ void VariableLookupAnalysis::dfs(const ExprLambda &Lambda,
     // Function arg cannot duplicate to it's formal.
     // If it this unluckily happens, we would like to skip this definition.
     if (!Arg.formals() || !Arg.formals()->dedup().contains(Arg.id()->name()))
-      DBuilder.add(Arg.id()->name(), Arg.id());
+      ToDef.insert_or_assign(Arg.id(),
+                             DBuilder.add(Arg.id()->name(), Arg.id()));
   }
 
   // { foo, bar, ... } : body
   ///  ^~~~~~~~~<--------------  add function formals.
   if (Arg.formals())
     for (const auto &[Name, Formal] : Arg.formals()->dedup())
-      DBuilder.add(Name, Formal->id());
+      ToDef.insert_or_assign(Formal->id(), DBuilder.add(Name, Formal->id()));
 
   auto NewEnv = std::make_shared<EnvNode>(Env, DBuilder.finish(), &Lambda);
 
@@ -136,7 +144,7 @@ VariableLookupAnalysis::dfsAttrs(const SemaAttrs &SA,
     DefBuilder DB;
     // For each static names, create a name binding.
     for (const auto &[Name, Attr] : SA.staticAttrs())
-      DB.add(Name, &Attr.key());
+      ToDef.insert_or_assign(&Attr.key(), DB.add(Name, &Attr.key()));
 
     auto NewEnv = std::make_shared<EnvNode>(Env, DB.finish(), Syntax);
 
@@ -205,8 +213,11 @@ void VariableLookupAnalysis::trivialDispatch(
 void VariableLookupAnalysis::dfs(const ExprWith &With,
                                  const std::shared_ptr<EnvNode> &Env) {
   auto NewEnv = std::make_shared<EnvNode>(Env, EnvNode::DefMap{}, &With);
-  if (!WithDefs.contains(&With))
-    WithDefs.insert({&With, std::make_shared<Definition>(&With.kwWith())});
+  if (!WithDefs.contains(&With)) {
+    auto NewDef = std::make_shared<Definition>(&With.kwWith());
+    ToDef.insert_or_assign(&With.kwWith(), NewDef);
+    WithDefs.insert_or_assign(&With, NewDef);
+  }
 
   if (With.with())
     dfs(*With.with(), Env);

libnixf/test/Sema/VariableLookup.cpp

@@ -3,6 +3,7 @@
 #include "nixf/Basic/Diagnostic.h"
 #include "nixf/Basic/Nodes/Expr.h"
 #include "nixf/Parse/Parser.h"
+#include "nixf/Sema/ParentMap.h"
 #include "nixf/Sema/VariableLookup.h"
 
 using namespace nixf;
@@ -171,4 +172,59 @@ TEST_F(VLATest, LivenessDupSymbol) {
   ASSERT_EQ(Diags[0].tags().size(), 0);
 }
 
+TEST_F(VLATest, ToDefAttrs) {
+  std::shared_ptr<Node> AST = parse("rec { x = 1; y = x; z = x; }", Diags);
+  VariableLookupAnalysis VLA(Diags);
+  VLA.runOnAST(*AST);
+
+  ParentMapAnalysis PMA;
+  PMA.runOnAST(*AST);
+
+  ASSERT_EQ(Diags.size(), 0);
+
+  ASSERT_TRUE(AST);
+  const Node *ID = AST->descend({{0, 6}, {0, 6}});
+  ID = PMA.upTo(*ID, Node::NK_AttrName);
+  ASSERT_EQ(ID->kind(), Node::NK_AttrName);
+
+  const auto *Def = VLA.toDef(*ID);
+  ASSERT_TRUE(Def);
+
+  ASSERT_EQ(Def->uses().size(), 2);
+}
+
+TEST_F(VLATest, ToDefLambda) {
+  std::shared_ptr<Node> AST =
+      parse("arg: { foo, bar} : arg + foo + bar", Diags);
+  VariableLookupAnalysis VLA(Diags);
+  VLA.runOnAST(*AST);
+
+  ASSERT_EQ(Diags.size(), 0);
+
+  ASSERT_TRUE(AST);
+  const Node *ID = AST->descend({{0, 1}, {0, 1}});
+  ASSERT_EQ(ID->kind(), Node::NK_Identifer);
+
+  const auto *Def = VLA.toDef(*ID);
+  ASSERT_TRUE(Def);
+
+  ASSERT_EQ(Def->uses().size(), 1);
+}
+
+TEST_F(VLATest, ToDefWith) {
+  std::shared_ptr<Node> AST = parse("with builtins; [ x y z ]", Diags);
+  VariableLookupAnalysis VLA(Diags);
+  VLA.runOnAST(*AST);
+
+  ASSERT_EQ(Diags.size(), 0);
+
+  ASSERT_TRUE(AST);
+  const Node *KwWith = AST->descend({{0, 1}, {0, 1}});
+
+  const auto *Def = VLA.toDef(*KwWith);
+  ASSERT_TRUE(Def);
+
+  ASSERT_EQ(Def->uses().size(), 3);
+}
+
 } // namespace

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