chore: Improve AST traversal in the LSP server

src/frontend/cxx/cxx_document.cc

@@ -21,13 +21,8 @@
 #include "cxx_document.h"
 
 #include <cxx/ast.h>
-#include <cxx/ast_slot.h>
-#include <cxx/ast_visitor.h>
 #include <cxx/control.h>
 #include <cxx/gcc_linux_toolchain.h>
-#include <cxx/lexer.h>
-#include <cxx/lsp/enums.h>
-#include <cxx/lsp/requests.h>
 #include <cxx/lsp/types.h>
 #include <cxx/macos_toolchain.h>
 #include <cxx/preprocessor.h>
@@ -209,43 +204,28 @@ auto CxxDocument::translationUnit() const -> TranslationUnit* {
   return &d->unit;
 }
 
-namespace {
-
-struct Visit {
-  CxxDocument* doc_;
-  std::function<bool(AST*)> visitor;
-  ASTSlot slotInfo;
-
-  void preVisit(AST* ast) {
-    if (!ast) return;
-
-    if (visitor(ast)) return;
-
-    // do a pre-visit using the low-level AST API
+auto CxxDocument::textOf(AST* ast) -> std::optional<std::string_view> {
+  return textInRange(ast->firstSourceLocation(), ast->lastSourceLocation());
+}
 
-    const auto slotCount = slotInfo(ast, 0).slotCount;
+auto CxxDocument::textInRange(SourceLocation start, SourceLocation end)
+    -> std::optional<std::string_view> {
+  auto& unit = d->unit;
+  auto preprocessor = unit.preprocessor();
 
-    for (int index = 0; index < slotCount; ++index) {
-      const auto childInfo = slotInfo(ast, index);
+  const auto startToken = unit.tokenAt(start);
+  const auto endToken = unit.tokenAt(end.previous());
 
-      if (childInfo.kind == ASTSlotKind::kNode) {
-        auto child = reinterpret_cast<AST*>(childInfo.handle);
-        if (child) preVisit(child);
-      } else if (childInfo.kind == ASTSlotKind::kNodeList) {
-        auto list = reinterpret_cast<List<AST*>*>(childInfo.handle);
-        for (auto node : ListView{list}) {
-          preVisit(node);
-        }
-      }
-    }
+  if (startToken.fileId() != endToken.fileId()) {
+    return std::nullopt;
   }
-};
 
-}  // namespace
+  std::string_view source = preprocessor->source(startToken.fileId());
+
+  const auto offset = startToken.offset();
+  const auto length = endToken.offset() + endToken.length() - offset;
 
-void CxxDocument::preVisit(std::function<bool(AST*)> visitor) {
-  auto ast = d->unit.ast();
-  Visit{this, std::move(visitor)}.preVisit(ast);
+  return source.substr(offset, length);
 }
 
 }  // namespace cxx::lsp

src/frontend/cxx/cxx_document.h

@@ -23,8 +23,7 @@
 #include <cxx/lsp/fwd.h>
 #include <cxx/translation_unit.h>
 
-#include <functional>
-#include <memory>
+#include <optional>
 #include <string>
 
 #include "cli.h"
@@ -43,7 +42,10 @@ class CxxDocument {
 
   [[nodiscard]] auto translationUnit() const -> TranslationUnit*;
 
-  void preVisit(std::function<bool(AST*)> visitor);
+  [[nodiscard]] auto textOf(AST* ast) -> std::optional<std::string_view>;
+
+  [[nodiscard]] auto textInRange(SourceLocation start, SourceLocation end)
+      -> std::optional<std::string_view>;
 
  private:
   struct Private;

src/frontend/cxx/lsp_server.cc

@@ -20,9 +20,14 @@
 
 #include "lsp_server.h"
 
+#include <cxx/ast.h>
 #include <cxx/lsp/enums.h>
 #include <cxx/lsp/requests.h>
 #include <cxx/lsp/types.h>
+#include <cxx/name_printer.h>
+#include <cxx/preprocessor.h>
+#include <cxx/symbols.h>
+#include <cxx/type_printer.h>
 #include <utf8/unchecked.h>
 
 #include <format>
@@ -433,12 +438,15 @@ void Server::operator()(DocumentSymbolRequest request) {
 
   auto uri = request.params().textDocument().uri();
   auto doc = latestDocument(uri);
+  auto id = request.id();
 
-  withUnsafeJson([&](json storage) {
-    DocumentSymbolResponse response(storage);
-    response.id(request.id());
-    (void)response.result();
-    sendToClient(response);
+  run([=, this] {
+    withUnsafeJson([&](json storage) {
+      DocumentSymbolResponse response(storage);
+      response.id(id);
+      (void)response.result();
+      sendToClient(response);
+    });
   });
 }
 
@@ -462,7 +470,7 @@ void Server::operator()(SetTraceNotification notification) {
   trace_ = notification.params().value();
 
   if (trace_ != TraceValue::kOff) {
-    logTrace("Trace level set to {}", to_string(trace_));
+    logTrace(std::format("Trace level set to {}", to_string(trace_)));
     return;
   }
 }

src/parser/cxx/ast.h

@@ -2262,6 +2262,7 @@ class InitDeclaratorAST final : public AST {
   DeclaratorAST* declarator = nullptr;
   RequiresClauseAST* requiresClause = nullptr;
   ExpressionAST* initializer = nullptr;
+  Symbol* symbol = nullptr;
 
   void accept(ASTVisitor* visitor) override { visitor->visit(this); }
 

src/parser/cxx/ast_cursor.cc

@@ -0,0 +1,83 @@
+// Copyright (c) 2024 Roberto Raggi <roberto.raggi@gmail.com>
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+// SOFTWARE.
+
+#include <cxx/ast.h>
+#include <cxx/ast_cursor.h>
+#include <cxx/ast_slot.h>
+
+namespace cxx {
+
+ASTCursor::ASTCursor(AST* root, std::string_view name) {
+  stack_.push_back(Node{root, name});
+}
+
+ASTCursor::~ASTCursor() {}
+
+void ASTCursor::step() {
+  ASTSlot slotInfo;
+
+  while (!stack_.empty()) {
+    auto [node, kind] = stack_.back();
+    stack_.pop_back();
+
+    if (auto ast = std::get_if<AST*>(&node)) {
+      if (!*ast) continue;
+
+      const auto slotCount = slotInfo(*ast, 0).slotCount;
+
+      std::vector<ASTSlot::SlotInfo> slots;
+      slots.reserve(slotCount);
+
+      for (int i = 0; i < slotCount; ++i) {
+        auto childInfo = slotInfo(*ast, i);
+        slots.push_back(childInfo);
+      }
+
+      for (const auto& slot : slots | std::views::reverse) {
+        auto name = to_string(slot.nameIndex);
+        if (slot.kind == ASTSlotKind::kNode) {
+          auto node = reinterpret_cast<AST*>(slot.handle);
+          stack_.push_back(Node{node, name});
+        } else if (slot.kind == ASTSlotKind::kNodeList) {
+          auto list = reinterpret_cast<List<AST*>*>(slot.handle);
+          stack_.push_back(Node{list, name});
+        }
+      }
+
+      break;
+    } else if (auto list = std::get_if<List<AST*>*>(&node)) {
+      if (!*list) continue;
+
+      std::vector<AST*> children;
+
+      for (auto ast : ListView{*list}) {
+        if (ast) children.push_back(ast);
+      }
+
+      for (const auto& ast : children | std::views::reverse) {
+        stack_.push_back(Node{ast, kind});
+      }
+
+      break;
+    }
+  }
+}
+
+}  // namespace cxx

src/parser/cxx/ast_cursor.h

@@ -0,0 +1,66 @@
+// Copyright (c) 2024 Roberto Raggi <roberto.raggi@gmail.com>
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+// SOFTWARE.
+
+#pragma once
+
+#include <cxx/ast_fwd.h>
+
+#include <deque>
+#include <string_view>
+#include <variant>
+
+namespace cxx {
+
+class ASTCursor {
+ public:
+  struct Node {
+    std::variant<AST*, List<AST*>*> node;
+    std::string_view name;
+  };
+
+  ASTCursor() = default;
+  ~ASTCursor();
+
+  ASTCursor(const ASTCursor&) = default;
+  ASTCursor& operator=(const ASTCursor&) = default;
+
+  ASTCursor(ASTCursor&&) = default;
+  ASTCursor& operator=(ASTCursor&&) = default;
+
+  ASTCursor(AST* root, std::string_view name);
+
+  explicit operator bool() const { return !empty(); }
+
+  [[nodiscard]] auto empty() const -> bool { return stack_.empty(); }
+
+  [[nodiscard]] auto operator*() const -> const Node& { return stack_.back(); }
+
+  void step();
+
+  auto operator++() -> ASTCursor& {
+    step();
+    return *this;
+  }
+
+ private:
+  std::deque<Node> stack_;
+};
+
+}  // namespace cxx

src/parser/cxx/parser.cc

@@ -14,9 +14,9 @@
 // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-// SOFTWARE.
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+// FRnewOM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+// DEALINGS IN THE SOFTWARE.
 
 #include <cxx/parser.h>
 
@@ -6726,6 +6726,7 @@ auto Parser::parse_init_declarator(InitDeclaratorAST*& yyast,
 auto Parser::parse_init_declarator(InitDeclaratorAST*& yyast,
                                    DeclaratorAST* declarator, Decl& decl)
     -> bool {
+  Symbol* declaredSynbol = nullptr;
   if (auto declId = decl.declaratorId; declId) {
     auto symbolType = GetDeclaratorType{this}(declarator, decl.specs.getType());
     const auto name = convertName(declId->unqualifiedId);
@@ -6735,22 +6736,26 @@ auto Parser::parse_init_declarator(InitDeclaratorAST*& yyast,
         symbol->setName(name);
         symbol->setType(symbolType);
         std::invoke(DeclareSymbol{this, scope_}, symbol);
+        declaredSynbol = symbol;
       } else if (getFunctionPrototype(declarator)) {
         auto functionSymbol =
             control_->newFunctionSymbol(scope_, decl.location());
         applySpecifiers(functionSymbol, decl.specs);
         functionSymbol->setName(name);
         functionSymbol->setType(symbolType);
         std::invoke(DeclareSymbol{this, scope_}, functionSymbol);
+        declaredSynbol = functionSymbol;
       } else {
         auto symbol = control_->newVariableSymbol(scope_, decl.location());
         applySpecifiers(symbol, decl.specs);
         symbol->setName(name);
         symbol->setType(symbolType);
         std::invoke(DeclareSymbol{this, scope_}, symbol);
+        declaredSynbol = symbol;
       }
     }
   }
+
   RequiresClauseAST* requiresClause = nullptr;
   ExpressionAST* initializer = nullptr;
 
@@ -6765,6 +6770,7 @@ auto Parser::parse_init_declarator(InitDeclaratorAST*& yyast,
   ast->declarator = declarator;
   ast->requiresClause = requiresClause;
   ast->initializer = initializer;
+  ast->symbol = declaredSynbol;
 
   return true;
 }

src/parser/cxx/preprocessor.cc

@@ -2796,6 +2796,17 @@ void Preprocessor::setOnWillIncludeHeader(
 
 void Preprocessor::squeeze() { d->pool_.reset(); }
 
+auto Preprocessor::sourceFileName(uint32_t sourceFileId) const
+    -> const std::string & {
+  assert(sourceFileId > 0);
+  return d->sourceFiles_[sourceFileId - 1]->fileName;
+}
+
+auto Preprocessor::source(uint32_t sourceFileId) const -> const std::string & {
+  assert(sourceFileId > 0);
+  return d->sourceFiles_[sourceFileId - 1]->source;
+}
+
 void Preprocessor::preprocess(std::string source, std::string fileName,
                               std::vector<Token> &tokens) {
   struct {

src/parser/cxx/preprocessor.h

@@ -119,6 +119,11 @@ class Preprocessor {
   [[nodiscard]] auto continuePreprocessing(std::vector<Token> &outputTokens)
       -> Status;
 
+  [[nodiscard]] auto sourceFileName(uint32_t sourceFileId) const
+      -> const std::string &;
+
+  [[nodiscard]] auto source(uint32_t sourceFileId) const -> const std::string &;
+
   void preprocess(std::string source, std::string fileName,
                   std::vector<Token> &tokens);