From 0fddb56b03d098d45a60742c9b6f8dedb251cf7d Mon Sep 17 00:00:00 2001
From: ef3d0c3e <ef3d0c3e@pundalik.org>
Date: Thu, 5 Dec 2024 09:40:00 +0100
Subject: [PATCH] Add conceal support

---
 src/elements/list.rs |  28 +++++++--
 src/lsp/conceal.rs   | 138 +++++++++++++++++++++++++++++++++++++++++++
 src/lsp/data.rs      |   7 +++
 src/lsp/hints.rs     |   1 +
 src/lsp/mod.rs       |   1 +
 src/lsp/semantic.rs  |   1 +
 src/server.rs        |  44 +++++++++++++-
 7 files changed, 214 insertions(+), 6 deletions(-)
 create mode 100644 src/lsp/conceal.rs

diff --git a/src/elements/list.rs b/src/elements/list.rs
index b24d8a5..47d672c 100644
--- a/src/elements/list.rs
+++ b/src/elements/list.rs
@@ -23,6 +23,7 @@ use crate::parser::source::Cursor;
 use crate::parser::source::Token;
 use crate::parser::source::VirtualSource;
 use crate::parser::util;
+use lsp::conceal::Conceals;
 use lsp::hints::Hints;
 use parser::util::escape_source;
 use regex::Regex;
@@ -330,6 +331,7 @@ impl Rule for ListRule {
 					offset.unwrap_or(usize::MAX),
 				);
 
+				// Semantic
 				if let Some((sems, tokens)) =
 					Semantics::from_source(cursor.source.clone(), &state.shared.lsp)
 				{
@@ -340,12 +342,30 @@ impl Rule for ListRule {
 					}
 				}
 
-				if let Some(hints) =
-					Hints::from_source(cursor.source.clone(), &state.shared.lsp)
+				if let Some(conceals) =
+					Conceals::from_source(cursor.source.clone(), &state.shared.lsp)
 				{
+					let mut i = captures.get(1).unwrap().start();
+					for (numbered, _) in &depth {
+						conceals.add(
+							i..i + 1,
+							lsp::conceal::ConcealTarget::Highlight {
+								text: if *numbered {
+									"⦾".into()
+								} else {
+									"⦿".into()
+								},
+								highlight_group: "Function".into(),
+							},
+						);
+						i += 1;
+					}
+				}
+
+				// Hints
+				if let Some(hints) = Hints::from_source(cursor.source.clone(), &state.shared.lsp) {
 					let mut label = String::new();
-					for (_, id) in &depth
-					{
+					for (_, id) in &depth {
 						if !label.is_empty() {
 							label.push('.');
 						}
diff --git a/src/lsp/conceal.rs b/src/lsp/conceal.rs
new file mode 100644
index 0000000..2ae0363
--- /dev/null
+++ b/src/lsp/conceal.rs
@@ -0,0 +1,138 @@
+use std::cell::Ref;
+use std::cell::RefCell;
+use std::ops::Range;
+use std::rc::Rc;
+
+use serde::Deserialize;
+use serde::Serialize;
+use tower_lsp::jsonrpc::Request;
+use tower_lsp::jsonrpc::{self};
+use tower_lsp::lsp_types::Position;
+
+use crate::parser::source::LineCursor;
+use crate::parser::source::Source;
+use crate::parser::source::SourceFile;
+use crate::parser::source::SourcePosition;
+use crate::parser::source::VirtualSource;
+
+use super::data::LSPData;
+
+#[derive(Debug, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct ConcealParams {
+	pub text_document: tower_lsp::lsp_types::TextDocumentIdentifier,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct ConcealInfo {
+	pub range: tower_lsp::lsp_types::Range,
+	pub conceal_text: ConcealTarget,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub enum ConcealTarget {
+	Text(String),
+	Highlight {
+		text: String,
+		highlight_group: String,
+	},
+}
+
+/// Per file conceals
+#[derive(Debug)]
+pub struct ConcealsData {
+	/// The current cursor
+	cursor: RefCell<LineCursor>,
+
+	/// The conceals
+	pub conceals: RefCell<Vec<ConcealInfo>>,
+}
+
+impl ConcealsData {
+	pub fn new(source: Rc<dyn Source>) -> Self {
+		Self {
+			cursor: RefCell::new(LineCursor::new(source)),
+			conceals: RefCell::new(vec![]),
+		}
+	}
+}
+
+/// Temporary data returned by [`Self::from_source_impl`]
+#[derive(Debug)]
+pub struct Conceals<'a> {
+	pub(self) conceals: Ref<'a, ConcealsData>,
+	// The source used when resolving the parent source
+	pub(self) original_source: Rc<dyn Source>,
+	/// The resolved parent source
+	pub(self) source: Rc<dyn Source>,
+}
+
+impl<'a> Conceals<'a> {
+	fn from_source_impl(
+		source: Rc<dyn Source>,
+		lsp: &'a Option<RefCell<LSPData>>,
+		original_source: Rc<dyn Source>,
+	) -> Option<Self> {
+		if (source.name().starts_with(":LUA:") || source.name().starts_with(":VAR:"))
+			&& source.downcast_ref::<VirtualSource>().is_some()
+		{
+			return None;
+		}
+
+		if let Some(location) = source
+			.clone()
+			.downcast_rc::<VirtualSource>()
+			.ok()
+			.as_ref()
+			.map(|parent| parent.location())
+			.unwrap_or(None)
+		{
+			return Self::from_source_impl(location.source(), lsp, original_source);
+		} else if let Ok(source) = source.clone().downcast_rc::<SourceFile>() {
+			return Ref::filter_map(lsp.as_ref().unwrap().borrow(), |lsp: &LSPData| {
+				lsp.conceals.get(&(source.clone() as Rc<dyn Source>))
+			})
+			.ok()
+			.map(|conceals| Self {
+				conceals,
+				source,
+				original_source,
+			});
+		}
+		None
+	}
+
+	pub fn from_source(source: Rc<dyn Source>, lsp: &'a Option<RefCell<LSPData>>) -> Option<Self> {
+		if lsp.is_none() {
+			return None;
+		}
+		Self::from_source_impl(source.clone(), lsp, source)
+	}
+
+	pub fn add(&self, range: Range<usize>, text: ConcealTarget) {
+		let range = self.original_source.original_range(range).1;
+		let mut cursor = self.conceals.cursor.borrow_mut();
+		cursor.move_to(range.start);
+
+		let line = cursor.line;
+		let start_char = cursor.line_pos;
+
+		cursor.move_to(range.end);
+		assert_eq!(line, cursor.line);
+		let end_char = cursor.line_pos;
+
+		self.conceals.conceals.borrow_mut().push(ConcealInfo {
+			range: tower_lsp::lsp_types::Range {
+				start: Position {
+					line: line as u32,
+					character: start_char as u32,
+				},
+				end: Position {
+					line: line as u32,
+					character: end_char as u32,
+				},
+			},
+			conceal_text: text,
+		})
+	}
+}
diff --git a/src/lsp/data.rs b/src/lsp/data.rs
index 32847c7..c933620 100644
--- a/src/lsp/data.rs
+++ b/src/lsp/data.rs
@@ -3,6 +3,7 @@ use std::rc::Rc;
 
 use crate::parser::source::Source;
 
+use super::conceal::ConcealsData;
 use super::definition::DefinitionData;
 use super::hints::HintsData;
 use super::semantic::SemanticsData;
@@ -14,6 +15,7 @@ pub struct LSPData {
 	pub semantic_data: HashMap<Rc<dyn Source>, SemanticsData>,
 	pub inlay_hints: HashMap<Rc<dyn Source>, HintsData>,
 	pub definitions: HashMap<Rc<dyn Source>, DefinitionData>,
+	pub conceals: HashMap<Rc<dyn Source>, ConcealsData>,
 }
 
 impl LSPData {
@@ -23,6 +25,7 @@ impl LSPData {
 			semantic_data: HashMap::new(),
 			inlay_hints: HashMap::new(),
 			definitions: HashMap::new(),
+			conceals: HashMap::new(),
 		}
 	}
 
@@ -40,5 +43,9 @@ impl LSPData {
 			self.definitions
 				.insert(source.clone(), DefinitionData::new());
 		}
+		if !self.conceals.contains_key(&source) {
+			self.conceals
+				.insert(source.clone(), ConcealsData::new(source.clone()));
+		}
 	}
 }
diff --git a/src/lsp/hints.rs b/src/lsp/hints.rs
index 0ce3e9e..b6806db 100644
--- a/src/lsp/hints.rs
+++ b/src/lsp/hints.rs
@@ -31,6 +31,7 @@ impl HintsData {
 	}
 }
 
+/// Temporary data returned by [`Self::from_source_impl`]
 #[derive(Debug)]
 pub struct Hints<'a> {
 	pub(self) hints: Ref<'a, HintsData>,
diff --git a/src/lsp/mod.rs b/src/lsp/mod.rs
index ba6ed88..082719d 100644
--- a/src/lsp/mod.rs
+++ b/src/lsp/mod.rs
@@ -2,3 +2,4 @@ pub mod data;
 pub mod definition;
 pub mod hints;
 pub mod semantic;
+pub mod conceal;
diff --git a/src/lsp/semantic.rs b/src/lsp/semantic.rs
index de93d07..1d86631 100644
--- a/src/lsp/semantic.rs
+++ b/src/lsp/semantic.rs
@@ -305,6 +305,7 @@ impl SemanticsData {
 	}
 }
 
+/// Temporary data returned by [`Self::from_source_impl`]
 #[derive(Debug)]
 pub struct Semantics<'a> {
 	pub(self) sems: Ref<'a, SemanticsData>,
diff --git a/src/server.rs b/src/server.rs
index 036e9ad..1dea4c6 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -9,12 +9,17 @@ mod parser;
 use std::rc::Rc;
 
 use dashmap::DashMap;
+use downcast_rs::Downcast;
+use lsp::conceal::ConcealInfo;
+use lsp::conceal::ConcealParams;
+use lsp::conceal::ConcealTarget;
 use parser::langparser::LangParser;
 use parser::parser::ParseMode;
 use parser::parser::Parser;
 use parser::parser::ParserState;
 use parser::reports::Report;
 use parser::source::SourceFile;
+use tower_lsp::jsonrpc;
 use tower_lsp::lsp_types::*;
 use tower_lsp::Client;
 use tower_lsp::LanguageServer;
@@ -29,6 +34,7 @@ struct Backend {
 	semantic_token_map: DashMap<String, Vec<SemanticToken>>,
 	diagnostic_map: DashMap<String, Vec<Diagnostic>>,
 	hints_map: DashMap<String, Vec<InlayHint>>,
+	conceals_map: DashMap<String, Vec<ConcealInfo>>,
 }
 
 #[derive(Debug)]
@@ -112,6 +118,36 @@ impl Backend {
 				}
 			}
 		}
+
+		// Conceals
+		if let Some(lsp) = state.shared.lsp.as_ref() {
+			let borrow = lsp.borrow();
+			for (source, conceals) in &borrow.conceals {
+				if let Some(path) = source
+					.clone()
+					.downcast_rc::<SourceFile>()
+					.ok()
+					.map(|source| source.path().to_owned())
+				{
+					self.conceals_map
+						.insert(path, conceals.conceals.replace(vec![]));
+				}
+			}
+		}
+	}
+
+	async fn handle_conceal_request(
+		&self,
+		params: ConcealParams,
+	) -> jsonrpc::Result<Vec<ConcealInfo>> {
+		eprintln!("HERE {:#?}", self.conceals_map);
+		if let Some(conceals) = self.conceals_map.get(params.text_document.uri.as_str()) {
+			let (_, data) = conceals.pair();
+
+		eprintln!("HERE2");
+			return Ok(data.to_vec());
+		}
+		Ok(vec![])
 	}
 }
 
@@ -313,13 +349,17 @@ async fn main() {
 	let stdin = tokio::io::stdin();
 	let stdout = tokio::io::stdout();
 
-	let (service, socket) = LspService::new(|client| Backend {
+	let (service, socket) = LspService::build(|client| Backend {
 		client,
 		document_map: DashMap::new(),
 		definition_map: DashMap::new(),
 		semantic_token_map: DashMap::new(),
 		diagnostic_map: DashMap::new(),
 		hints_map: DashMap::new(),
-	});
+		conceals_map: DashMap::new(),
+	})
+	.custom_method("textDocument/conceal", Backend::handle_conceal_request)
+	.finish();
+
 	Server::new(stdin, stdout, socket).serve(service).await;
 }