From 9214c937581b667bfb436e3221c48667f25c1dde Mon Sep 17 00:00:00 2001
From: Yingchi Long <i@lyc.dev>
Date: Tue, 23 Jul 2024 16:49:10 +0800
Subject: [PATCH] random

---
 nixd/lib/Controller/Definition.cpp | 77 +++++++++++++++++++++++-------
 1 file changed, 61 insertions(+), 16 deletions(-)

diff --git a/nixd/lib/Controller/Definition.cpp b/nixd/lib/Controller/Definition.cpp
index 829ca59d3..3cd227e08 100644
--- a/nixd/lib/Controller/Definition.cpp
+++ b/nixd/lib/Controller/Definition.cpp
@@ -10,6 +10,8 @@
 #include "nixd/Controller/Controller.h"
 #include "nixd/Protocol/AttrSet.h"
 
+#include "lspserver/Protocol.h"
+
 #include <boost/asio/post.hpp>
 
 #include <llvm/Support/Error.h>
@@ -94,6 +96,13 @@ const ExprVar *findVar(const Node &N, const ParentMapAnalysis &PMA,
   return static_cast<const ExprVar *>(PMA.upTo(N, Node::NK_ExprVar));
 }
 
+Location defToLocaton(const Definition &Def, URIForFile URI) {
+  return Location{
+      .uri = std::move(URI),
+      .range = toLSPRange(Def.syntax()->range()),
+  };
+}
+
 Expected<Location> staticDef(URIForFile URI, const Node &N,
                              const ParentMapAnalysis &PMA,
                              const VariableLookupAnalysis &VLA) {
@@ -104,11 +113,6 @@ Expected<Location> staticDef(URIForFile URI, const Node &N,
   if (ExpDef->isBuiltin())
     return error("this is a builtin variable defined by nix interpreter");
   assert(ExpDef->syntax());
-
-  return Location{
-      .uri = std::move(URI),
-      .range = toLSPRange(ExpDef->syntax()->range()),
-  };
 }
 
 struct NoLocationsFoundInNixpkgsException : std::exception {
@@ -270,6 +274,38 @@ Locations defineSelect(const ExprSelect &Sel, const VariableLookupAnalysis &VLA,
   return {};
 }
 
+Locations defineVarStatic(const ExprVar &Var, const VariableLookupAnalysis &VLA,
+                          const ParentMapAnalysis &PM, const URIForFile &URI) {
+  Expected<Location> StaticLoc = staticDef(URI, Var, PM, VLA);
+  if (!StaticLoc) {
+    elog("definition/static: {0}", StaticLoc.takeError());
+    return {*StaticLoc};
+  }
+  return {};
+}
+
+template <class T>
+std::vector<T> mergeVec(std::vector<T> A, const std::vector<T> &B) {
+  A.insert(A.end(), B.begin(), B.end());
+  return A;
+}
+
+Locations defineVar(const ExprVar &Var, const VariableLookupAnalysis &VLA,
+                    const ParentMapAnalysis &PM, AttrSetClient &NixpkgsClient,
+                    const URIForFile &URI) {
+  Locations StaticLocs = defineVarStatic(Var, VLA, PM, URI);
+
+  // Nixpkgs locations.
+  try {
+    Selector Sel = mkIdiomSelector(Var, VLA, PM);
+    Locations NixpkgsLocs = defineNixpkgsSelector(Sel, NixpkgsClient);
+    return mergeVec(std::move(StaticLocs), NixpkgsLocs);
+  } catch (IdiomSelectorException &E) {
+    elog("definition/idiom/selector: {0}", E.what());
+  }
+  return StaticLocs;
+}
+
 /// \brief Squash a vector into smaller json variant.
 template <class T> llvm::json::Value squash(std::vector<T> List) {
   std::size_t Size = List.size();
@@ -291,6 +327,21 @@ llvm::Expected<llvm::json::Value> squash(llvm::Expected<std::vector<T>> List) {
   return squash(std::move(*List));
 }
 
+const Definition &findVarDefinition(const ExprVar &Var,
+                                    const VariableLookupAnalysis &VLA) {
+  LookupResult Result = VLA.query(static_cast<const ExprVar &>(Var));
+
+  if (Result.Kind == ResultKind::Undefined)
+    return error("this varaible is undefined");
+
+  if (Result.Kind == ResultKind::NoSuchVar)
+    throw NoSuchVarException();
+
+  assert(Result.Def);
+
+  return *Result.Def;
+}
+
 } // namespace
 
 Expected<const Definition &>
@@ -303,17 +354,7 @@ nixd::findDefinition(const Node &N, const ParentMapAnalysis &PMA,
     return error("cannot find variable on given position");
   }
   assert(Var->kind() == Node::NK_ExprVar);
-  LookupResult Result = VLA.query(static_cast<const ExprVar &>(*Var));
-
-  if (Result.Kind == ResultKind::Undefined)
-    return error("this varaible is undefined");
-
-  if (Result.Kind == ResultKind::NoSuchVar)
-    return error("this varaible is not used in var lookup (duplicated attr?)");
-
-  assert(Result.Def);
-
-  return *Result.Def;
+  return findVarDefinition(*Var, VLA);
 }
 
 void Controller::onDefinition(const TextDocumentPositionParams &Params,
@@ -341,6 +382,10 @@ void Controller::onDefinition(const TextDocumentPositionParams &Params,
 
         return Reply(squash([&]() -> llvm::Expected<Locations> {
           switch (UpExpr.kind()) {
+          case Node::NK_ExprVar: {
+            const auto &Var = static_cast<const ExprVar &>(UpExpr);
+            return defineVar(Var, VLA, PM, *nixpkgsClient(), URI);
+          }
           case Node::NK_ExprSelect: {
             const auto &Sel = static_cast<const ExprSelect &>(UpExpr);
             return defineSelect(Sel, VLA, PM, *nixpkgsClient());