nixd/lib/Controller: use CheckReturn to simplify early-returns (NFC)

nixd/include/nixd/Controller/Controller.h

@@ -106,34 +106,33 @@ class Controller : public lspserver::LSPServer {
       EndWorkDoneProgress(Params);
   }
 
-  std::mutex TUsLock;
+  mutable std::mutex TUsLock;
   llvm::StringMap<std::shared_ptr<NixTU>> TUs;
 
-  template <class T>
-  std::shared_ptr<NixTU> getTU(std::string File,
-                               lspserver::Callback<T> &Reply) {
+  std::shared_ptr<const NixTU> getTU(std::string_view File) const {
     using lspserver::error;
     std::lock_guard G(TUsLock);
     if (!TUs.count(File)) [[unlikely]] {
-      Reply(T{}); // Reply a default constructed response.
       lspserver::elog("cannot get translation unit: {0}", File);
       return nullptr;
     }
-    return TUs[File];
+    return TUs.lookup(File);
   }
 
-  template <class T>
-  std::shared_ptr<nixf::Node> getAST(const NixTU &TU,
-                                     lspserver::Callback<T> &Reply) {
+  static std::shared_ptr<nixf::Node> getAST(const NixTU &TU) {
     using lspserver::error;
     if (!TU.ast()) {
-      Reply(T{});
       lspserver::elog("AST is null on this unit");
       return nullptr;
     }
     return TU.ast();
   }
 
+  std::shared_ptr<const nixf::Node> getAST(std::string_view File) const {
+    auto TU = getTU(File);
+    return TU ? getAST(*TU) : nullptr;
+  }
+
   boost::asio::thread_pool Pool;
 
   /// Action right after a document is added (including updates).
@@ -217,8 +216,9 @@ class Controller : public lspserver::LSPServer {
   void onRename(const lspserver::RenameParams &Params,
                 lspserver::Callback<lspserver::WorkspaceEdit> Reply);
 
-  void onPrepareRename(const lspserver::TextDocumentPositionParams &Params,
-                       lspserver::Callback<lspserver::Range> Reply);
+  void
+  onPrepareRename(const lspserver::TextDocumentPositionParams &Params,
+                  lspserver::Callback<std::optional<lspserver::Range>> Reply);
 
   void onFormat(const lspserver::DocumentFormattingParams &Params,
                 lspserver::Callback<std::vector<lspserver::TextEdit>> Reply);

nixd/lib/Controller/CheckReturn.h

@@ -0,0 +1,16 @@
+#pragma once
+
+/// \brief Used for simplify early-returns.
+///
+///   const auto *foo = checkReturn(get(), nullptr)
+#define CheckReturn(x, Ret)                                                    \
+  ({                                                                           \
+    decltype(x) temp = (x);                                                    \
+    if (!temp) {                                                               \
+      return Ret;                                                              \
+    }                                                                          \
+    temp;                                                                      \
+  })
+
+/// \brief Variant of `CheckReturn`, but returns default constructed `CheckTy`
+#define CheckDefault(x) CheckReturn(x, CheckTy{})

nixd/lib/Controller/CodeAction.cpp

@@ -3,6 +3,7 @@
 /// [Code Action]:
 /// https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_codeAction
 
+#include "CheckReturn.h"
 #include "Convert.h"
 
 #include "nixd/Controller/Controller.h"
@@ -16,12 +17,15 @@ using namespace lspserver;
 
 void Controller::onCodeAction(const lspserver::CodeActionParams &Params,
                               Callback<std::vector<CodeAction>> Reply) {
+  using CheckTy = std::vector<CodeAction>;
   std::string File(Params.textDocument.uri.file());
   Range Range = Params.range;
   auto Action = [Reply = std::move(Reply), File, Range, this]() mutable {
-    if (auto TU = getTU(File, Reply)) {
-      std::vector<nixf::Diagnostic> Diagnostics = TU->diagnostics();
-      std::vector<CodeAction> Actions;
+    return Reply([&]() -> llvm::Expected<CheckTy> {
+      const auto TU = CheckDefault(getTU(File));
+
+      const auto &Diagnostics = TU->diagnostics();
+      auto Actions = std::vector<CodeAction>();
       Actions.reserve(Diagnostics.size());
       for (const nixf::Diagnostic &D : Diagnostics) {
         auto DRange = toLSPRange(TU->src(), D.range());
@@ -50,8 +54,8 @@ void Controller::onCodeAction(const lspserver::CodeActionParams &Params,
           });
         }
       }
-      Reply(std::move(Actions));
-    }
+      return Actions;
+    }());
   };
   boost::asio::post(Pool, std::move(Action));
 }

nixd/lib/Controller/Completion.cpp

@@ -4,6 +4,7 @@
 /// https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_completion
 
 #include "AST.h"
+#include "CheckReturn.h"
 #include "Convert.h"
 
 #include "lspserver/Protocol.h"
@@ -376,72 +377,60 @@ void completeSelect(const nixf::ExprSelect &Select, AttrSetClient &Client,
 
 void Controller::onCompletion(const CompletionParams &Params,
                               Callback<CompletionList> Reply) {
+  using CheckTy = CompletionList;
   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;
-        }
-        if (!Desc->children().empty()) {
-          Reply(CompletionList{});
-          return;
-        }
-        const nixf::Node &N = *Desc;
-        const ParentMapAnalysis &PM = *TU->parentMap();
-        const Node *MaybeUpExpr = PM.upExpr(N);
-        if (!MaybeUpExpr) {
-          // If there is no concrete expression containing the cursor
-          // Reply an empty list.
-          Reply(CompletionList{});
-          return;
-        }
-        // Otherwise, construct the completion list from a set of providers.
-        const Node &UpExpr = *MaybeUpExpr;
-        Reply([&]() -> CompletionList {
-          CompletionList List;
-          const VariableLookupAnalysis &VLA = *TU->variableLookup();
-          try {
-            switch (UpExpr.kind()) {
-            // In these cases, assume the cursor have "variable" scoping.
-            case Node::NK_ExprVar: {
-              completeVarName(VLA, PM,
-                              static_cast<const nixf::ExprVar &>(UpExpr),
-                              *nixpkgsClient(), List.items);
-              break;
-            }
-            // A "select" expression. e.g.
-            // foo.a|
-            // foo.|
-            // foo.a.bar|
-            case Node::NK_ExprSelect: {
-              const auto &Select =
-                  static_cast<const nixf::ExprSelect &>(UpExpr);
-              completeSelect(Select, *nixpkgsClient(), VLA, PM,
-                             N.kind() == Node::NK_Dot, List.items);
-              break;
-            }
-            case Node::NK_ExprAttrs: {
-              completeAttrPath(N, PM, OptionsLock, Options,
-                               ClientCaps.CompletionSnippets, List.items);
-              break;
-            }
-            default:
-              break;
-            }
-          } catch (ExceedSizeError &Err) {
-            List.isIncomplete = true;
+    const auto File = URI.file().str();
+    return Reply([&]() -> llvm::Expected<CompletionList> {
+      const auto TU = CheckDefault(getTU(File));
+      const auto AST = CheckDefault(getAST(*TU));
+
+      const auto *Desc = AST->descend({Pos, Pos});
+      CheckDefault(Desc && Desc->children().empty());
+
+      const auto &N = *Desc;
+      const auto &PM = *TU->parentMap();
+      const auto &UpExpr = *CheckDefault(PM.upExpr(N));
+
+      return [&]() {
+        CompletionList List;
+        const VariableLookupAnalysis &VLA = *TU->variableLookup();
+        try {
+          switch (UpExpr.kind()) {
+          // In these cases, assume the cursor have "variable" scoping.
+          case Node::NK_ExprVar: {
+            completeVarName(VLA, PM, static_cast<const nixf::ExprVar &>(UpExpr),
+                            *nixpkgsClient(), List.items);
+            return List;
+          }
+          // A "select" expression. e.g.
+          // foo.a|
+          // foo.|
+          // foo.a.bar|
+          case Node::NK_ExprSelect: {
+            const auto &Select = static_cast<const nixf::ExprSelect &>(UpExpr);
+            completeSelect(Select, *nixpkgsClient(), VLA, PM,
+                           N.kind() == Node::NK_Dot, List.items);
+            return List;
+          }
+          case Node::NK_ExprAttrs: {
+            completeAttrPath(N, PM, OptionsLock, Options,
+                             ClientCaps.CompletionSnippets, List.items);
+            return List;
           }
+          default:
+            return List;
+          }
+        } catch (ExceedSizeError &Err) {
+          List.isIncomplete = true;
           return List;
-        }());
-      }
-    }
+        }
+      }();
+    }());
   };
   boost::asio::post(Pool, std::move(Action));
 }
+#undef check
 
 void Controller::onCompletionItemResolve(const CompletionItem &Params,
                                          Callback<CompletionItem> Reply) {

nixd/lib/Controller/Definition.cpp

@@ -5,6 +5,7 @@
 
 #include "Definition.h"
 #include "AST.h"
+#include "CheckReturn.h"
 #include "Convert.h"
 
 #include "nixd/Controller/Controller.h"
@@ -354,52 +355,43 @@ const Definition &nixd::findDefinition(const Node &N,
   return findVarDefinition(*Var, VLA);
 }
 
+#define Check(x) CheckReturn(x, nullptr)
 void Controller::onDefinition(const TextDocumentPositionParams &Params,
                               Callback<llvm::json::Value> 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<Node> AST = getAST(*TU, Reply)) [[likely]] {
-        const VariableLookupAnalysis &VLA = *TU->variableLookup();
-        const ParentMapAnalysis &PM = *TU->parentMap();
-        const Node *MaybeN = AST->descend({Pos, Pos});
-        if (!MaybeN) [[unlikely]] {
-          Reply(error("cannot find AST node on given position"));
-          return;
+    const auto File = URI.file().str();
+    return Reply([&]() -> llvm::Expected<llvm::json::Value> {
+      const auto TU = Check(getTU(File));
+      const auto AST = Check(getAST(*TU));
+      const auto &VLA = *TU->variableLookup();
+      const auto &PM = *TU->parentMap();
+      const auto &N = *Check(AST->descend({Pos, Pos}));
+      const auto &UpExpr = *Check(PM.upExpr(N));
+
+      return squash([&]() -> llvm::Expected<Locations> {
+        // Special case for inherited names.
+        if (const ExprVar *Var = findInheritVar(N, PM, VLA))
+          return defineVar(*Var, VLA, PM, *nixpkgsClient(), URI, TU->src());
+
+        switch (UpExpr.kind()) {
+        case Node::NK_ExprVar: {
+          const auto &Var = static_cast<const ExprVar &>(UpExpr);
+          return defineVar(Var, VLA, PM, *nixpkgsClient(), URI, TU->src());
         }
-        const Node &N = *MaybeN;
-        const Node *MaybeUpExpr = PM.upExpr(N);
-        if (!MaybeUpExpr) {
-          Reply(nullptr);
-          return;
+        case Node::NK_ExprSelect: {
+          const auto &Sel = static_cast<const ExprSelect &>(UpExpr);
+          return defineSelect(Sel, VLA, PM, *nixpkgsClient());
         }
-
-        const Node &UpExpr = *MaybeUpExpr;
-
-        return Reply(squash([&]() -> llvm::Expected<Locations> {
-          // Special case for inherited names.
-          if (const ExprVar *Var = findInheritVar(N, PM, VLA))
-            return defineVar(*Var, VLA, PM, *nixpkgsClient(), URI, TU->src());
-
-          switch (UpExpr.kind()) {
-          case Node::NK_ExprVar: {
-            const auto &Var = static_cast<const ExprVar &>(UpExpr);
-            return defineVar(Var, VLA, PM, *nixpkgsClient(), URI, TU->src());
-          }
-          case Node::NK_ExprSelect: {
-            const auto &Sel = static_cast<const ExprSelect &>(UpExpr);
-            return defineSelect(Sel, VLA, PM, *nixpkgsClient());
-          }
-          case Node::NK_ExprAttrs:
-            return defineAttrPath(N, PM, OptionsLock, Options);
-          default:
-            break;
-          }
-          return error("unknown node type for definition");
-        }()));
-      }
-    }
+        case Node::NK_ExprAttrs:
+          return defineAttrPath(N, PM, OptionsLock, Options);
+        default:
+          break;
+        }
+        return error("unknown node type for definition");
+      }());
+    }());
   };
   boost::asio::post(Pool, std::move(Action));
 }
+#undef Check

nixd/lib/Controller/DocumentHighlight.cpp

@@ -3,6 +3,7 @@
 /// [Document Highlight]:
 /// https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#textDocument_documentHighlight
 
+#include "CheckReturn.h"
 #include "Convert.h"
 #include "Definition.h"
 
@@ -55,26 +56,24 @@ std::vector<DocumentHighlight> highlight(const nixf::Node &Desc,
 void Controller::onDocumentHighlight(
     const TextDocumentPositionParams &Params,
     Callback<std::vector<DocumentHighlight>> Reply) {
+  using CheckTy = std::vector<DocumentHighlight>;
   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;
-        }
-        try {
-          const auto &PM = *TU->parentMap();
-          const auto &VLA = *TU->variableLookup();
-          return Reply(highlight(*Desc, PM, VLA, URI, TU->src()));
-        } catch (std::exception &E) {
-          elog("textDocument/documentHighlight failed: {0}", E.what());
-          return Reply(std::vector<DocumentHighlight>{});
-        }
+    return Reply([&]() -> llvm::Expected<CheckTy> {
+      const auto TU = CheckDefault(getTU(File));
+      const auto AST = CheckDefault(getAST(*TU));
+
+      const auto &Desc = *CheckDefault(AST->descend({Pos, Pos}));
+      try {
+        const auto &PM = *TU->parentMap();
+        const auto &VLA = *TU->variableLookup();
+        return highlight(Desc, PM, VLA, URI, TU->src());
+      } catch (std::exception &E) {
+        elog("textDocument/documentHighlight failed: {0}", E.what());
+        return CheckTy{};
       }
-    }
+    }());
   };
   boost::asio::post(Pool, std::move(Action));
 }