From 21ba624dcfdf51cd849ebdf788d26006b0efdae8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vid=20Drobni=C4=8D?= <me@viddrobnic.com>
Date: Mon, 19 Aug 2024 18:56:36 +0200
Subject: [PATCH] feat: document symbols support

---
 language_server/src/analyze/mod.rs            |  2 +-
 language_server/src/analyze/symbol_info.rs    |  2 +-
 language_server/src/lib.rs                    | 14 +++++++
 .../src/message/document_symbol.rs            | 40 +++++++++++++++++++
 language_server/src/message/initialize.rs     |  1 +
 language_server/src/message/mod.rs            |  1 +
 6 files changed, 58 insertions(+), 2 deletions(-)
 create mode 100644 language_server/src/message/document_symbol.rs

diff --git a/language_server/src/analyze/mod.rs b/language_server/src/analyze/mod.rs
index 84bbf92..24d9ca1 100644
--- a/language_server/src/analyze/mod.rs
+++ b/language_server/src/analyze/mod.rs
@@ -11,9 +11,9 @@ use symbol_table::SymbolTable;
 
 pub mod document_info;
 pub mod location;
+pub mod symbol_info;
 
 mod documentation;
-mod symbol_info;
 mod symbol_table;
 
 pub fn analyze(program: &ast::Program) -> DocumentInfo {
diff --git a/language_server/src/analyze/symbol_info.rs b/language_server/src/analyze/symbol_info.rs
index 423f25c..62fbdee 100644
--- a/language_server/src/analyze/symbol_info.rs
+++ b/language_server/src/analyze/symbol_info.rs
@@ -1,6 +1,6 @@
 use parser::position::Range;
 
-#[derive(Debug, PartialEq, Eq)]
+#[derive(Debug, PartialEq, Eq, Clone, Copy)]
 pub enum DocumentSymbolKind {
     Function = 12,
     Variable = 13,
diff --git a/language_server/src/lib.rs b/language_server/src/lib.rs
index da57681..c03f47a 100644
--- a/language_server/src/lib.rs
+++ b/language_server/src/lib.rs
@@ -8,6 +8,7 @@ use std::{
 
 use analyze::{analyze, document_info::DocumentInfo};
 use diagnostics::{Diagnostic, DiagnosticSeverity, PublishDiagnosticsParams};
+use document_symbol::{DocumentSymbol, DocumentSymbolParams};
 use error::{Error, ErrorKind};
 use hover::{Hover, MarkupContent, MarkupKind};
 use message::{initialize::*, *};
@@ -266,6 +267,18 @@ impl Server {
 
                 Response::new_ok(req_id, res)
             }
+            "textDocument/documentSymbol" => {
+                let (req_id, params) = req.extract::<DocumentSymbolParams>()?;
+                let doc_name = params.text_document.uri.clone();
+
+                let doc_info = self.documents.get(&doc_name);
+                let mut res: Option<Vec<DocumentSymbol>> = None;
+                if let Some(doc_info) = doc_info {
+                    res = Some(DocumentSymbol::map_tree(&doc_info.symbol_tree));
+                }
+
+                Response::new_ok(req_id, res)
+            }
 
             method => {
                 self.log(LogLevel::Warn, &format!("Got unknown method: {method}"));
@@ -295,6 +308,7 @@ impl Server {
                 document_highlight_provider: true,
                 references_provider: true,
                 hover_provider: true,
+                document_symbol_provider: true,
             },
         }
     }
diff --git a/language_server/src/message/document_symbol.rs b/language_server/src/message/document_symbol.rs
new file mode 100644
index 0000000..de65aa1
--- /dev/null
+++ b/language_server/src/message/document_symbol.rs
@@ -0,0 +1,40 @@
+use parser::position::Range;
+use serde::{Deserialize, Serialize};
+
+use crate::{analyze::symbol_info, TextDocumentIdentifier};
+
+#[derive(Debug, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct DocumentSymbolParams {
+    pub text_document: TextDocumentIdentifier,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct DocumentSymbol {
+    pub name: String,
+    pub kind: u8,
+    pub range: Range,
+    pub selection_range: Range,
+    pub children: Vec<DocumentSymbol>,
+}
+
+impl DocumentSymbol {
+    pub fn map_tree(tree: &[symbol_info::DocumentSymbol]) -> Vec<DocumentSymbol> {
+        tree.iter()
+            .filter_map(|sym| {
+                let Some(name) = &sym.name else {
+                    return None;
+                };
+
+                Some(DocumentSymbol {
+                    name: name.clone(),
+                    kind: sym.kind as u8,
+                    range: sym.range,
+                    selection_range: sym.name_range,
+                    children: DocumentSymbol::map_tree(&sym.children),
+                })
+            })
+            .collect()
+    }
+}
diff --git a/language_server/src/message/initialize.rs b/language_server/src/message/initialize.rs
index 86cb851..646df19 100644
--- a/language_server/src/message/initialize.rs
+++ b/language_server/src/message/initialize.rs
@@ -29,6 +29,7 @@ pub struct ServerCapabilities {
     pub document_highlight_provider: bool,
     pub references_provider: bool,
     pub hover_provider: bool,
+    pub document_symbol_provider: bool,
 }
 
 #[derive(Debug, Serialize, Deserialize)]
diff --git a/language_server/src/message/mod.rs b/language_server/src/message/mod.rs
index a32f6d5..fc0b11e 100644
--- a/language_server/src/message/mod.rs
+++ b/language_server/src/message/mod.rs
@@ -7,6 +7,7 @@ use serde_json::Value;
 use crate::error::{Error, ErrorKind};
 
 pub mod diagnostics;
+pub mod document_symbol;
 pub mod hover;
 pub mod initialize;
 pub mod reference;