From b305b7936f88e1caf927fa8edc8c1ea51e32c547 Mon Sep 17 00:00:00 2001
From: Felix Zeller <felixazeller@gmail.com>
Date: Wed, 27 Mar 2024 21:10:10 -0400
Subject: [PATCH] fix: semantic tokens performance

---
 src/main.rs   | 25 +++++++++++++++++++++++++
 src/tokens.rs | 15 ++++++++-------
 2 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index af04e0f7..3cb421fd 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,3 +1,4 @@
+use std::collections::btree_set::Intersection;
 use std::collections::HashSet;
 use std::ops::{Deref, DerefMut};
 use std::path::{Path, PathBuf};
@@ -149,6 +150,9 @@ impl Backend {
     }
 
     async fn publish_diagnostics(&self) -> Result<()> {
+
+        let timer = std::time::Instant::now();
+
         self.client
             .log_message(MessageType::WARNING, "Diagnostics Started")
             .await;
@@ -185,6 +189,16 @@ impl Backend {
             .log_message(MessageType::WARNING, "Diagnostics Done")
             .await;
 
+
+        let elapsed = timer.elapsed();
+
+        self.client
+            .log_message(
+                MessageType::WARNING,
+                format!("Diagnostics Done took {}ms", elapsed.as_millis()),
+            )
+            .await;
+
         Ok(())
     }
 
@@ -602,11 +616,22 @@ impl LanguageServer for Backend {
 
         let settings = self.bind_settings(|settings| Ok(settings.clone())).await?;
 
+        let timer = std::time::Instant::now();
+
         let path = params_path!(params)?;
         let res = self
             .bind_vault(|vault| Ok(tokens::semantic_tokens_full(vault, &path, params, &settings)))
             .await;
 
+        let elapsed = timer.elapsed();
+
+        self.client
+            .log_message(
+                MessageType::WARNING,
+                format!("Semantic Tokens Done took {}ms", elapsed.as_millis()),
+            )
+            .await;
+
         return res;
     }
 }
diff --git a/src/tokens.rs b/src/tokens.rs
index d9371227..85c4724b 100644
--- a/src/tokens.rs
+++ b/src/tokens.rs
@@ -1,9 +1,10 @@
-use std::{iter, path::Path};
+use std::{collections::HashSet, iter, path::Path};
 
 use itertools::Itertools;
+use rayon::iter::{IntoParallelIterator, ParallelIterator};
 use tower_lsp::lsp_types::{SemanticToken, SemanticTokensParams, SemanticTokensResult};
 
-use crate::{config::Settings, vault::{Referenceable, Vault}};
+use crate::{config::Settings, diagnostics::path_unresolved_references, vault::{Referenceable, Vault}};
 
 pub fn semantic_tokens_full(
     vault: &Vault,
@@ -17,6 +18,8 @@ pub fn semantic_tokens_full(
 
     let references_in_file = vault.select_references(Some(path))?;
 
+    let path_unresolved: Option<HashSet<_>> = path_unresolved_references(vault, path).map(|thing| thing.into_par_iter().map(|(_, reference)| reference).collect());
+
     let tokens = references_in_file
         .into_iter()
         .sorted_by_key(|(_, reference)| {
@@ -28,9 +31,7 @@ pub fn semantic_tokens_full(
         .fold(vec![], |acc, (path, reference)| {
             let range = reference.data().range;
 
-            let is_unresolved = vault.select_referenceables_for_reference(reference, path)
-                .into_iter()
-                .any(|referenceable| referenceable.is_unresolved());
+            let is_unresolved = path_unresolved.as_ref().is_some_and(|unresolved| unresolved.contains(reference));
 
             match acc[..] {
                 [] => vec![(
@@ -62,9 +63,9 @@ pub fn semantic_tokens_full(
                     .collect_vec(),
             }
         })
-        .into_iter()
+        .into_par_iter()
         .map(|(_, token)| token)
-        .collect_vec(); // TODO: holy this is bad
+        .collect::<Vec<_>>(); // TODO: holy this is bad
 
     Some(SemanticTokensResult::Tokens(
         tower_lsp::lsp_types::SemanticTokens {