From 7aa24394222ea94c5c110e18c92951456acc20ea Mon Sep 17 00:00:00 2001
From: Tad Hardesty <tad@platymuus.com>
Date: Sat, 4 Dec 2021 17:58:53 -0800
Subject: [PATCH] Populate find-all-references table in a background thread

This hovers around 0.57s on my system, which will now be deducted from
the total time to load the environment.
---
 crates/dm-langserver/src/background.rs | 60 ++++++++++++++++++++++++++
 crates/dm-langserver/src/main.rs       | 21 ++++++---
 2 files changed, 74 insertions(+), 7 deletions(-)
 create mode 100644 crates/dm-langserver/src/background.rs

diff --git a/crates/dm-langserver/src/background.rs b/crates/dm-langserver/src/background.rs
new file mode 100644
index 00000000..db04e008
--- /dev/null
+++ b/crates/dm-langserver/src/background.rs
@@ -0,0 +1,60 @@
+//! Helper for running background tasks.
+
+use std::sync::mpsc::{channel, Receiver, TryRecvError};
+use std::thread;
+
+pub struct Background<T> {
+    value: Option<T>,
+    rx: Option<Receiver<T>>,
+}
+
+impl<T> Default for Background<T> {
+    fn default() -> Self {
+        Background {
+            value: None,
+            rx: None,
+        }
+    }
+}
+
+#[allow(dead_code)]
+impl<T: Send + 'static> Background<T> {
+    pub fn new<F: FnOnce() -> T + Send + 'static>(f: F) -> Self {
+        let mut this = Self::default();
+        this.spawn(f);
+        this
+    }
+
+    pub fn spawn<F: FnOnce() -> T + Send + 'static>(&mut self, f: F) {
+        self.rx = Some(spawn(f));
+    }
+
+    pub fn poll(&mut self) -> &mut Self {
+        if let Some(rx) = self.rx.take() {
+            match rx.try_recv() {
+                Ok(v) => {
+                    self.value = Some(v);
+                }
+                Err(TryRecvError::Empty) => self.rx = Some(rx),
+                Err(TryRecvError::Disconnected) => {}
+            }
+        }
+        self
+    }
+
+    pub fn is_busy(&self) -> bool {
+        self.rx.is_some()
+    }
+
+    pub fn value(&self) -> Option<&T> {
+        self.value.as_ref()
+    }
+}
+
+fn spawn<R: Send + 'static, F: FnOnce() -> R + Send + 'static>(f: F) -> Receiver<R> {
+    let (tx, rx) = channel();
+    thread::spawn(move || {
+        let _ = tx.send(f());
+    });
+    rx
+}
diff --git a/crates/dm-langserver/src/main.rs b/crates/dm-langserver/src/main.rs
index c53980ec..e0ba57b8 100644
--- a/crates/dm-langserver/src/main.rs
+++ b/crates/dm-langserver/src/main.rs
@@ -29,6 +29,7 @@ mod find_references;
 mod extras;
 mod completion;
 mod color;
+mod background;
 
 mod debugger;
 
@@ -155,7 +156,7 @@ struct Engine<'a> {
     context: &'a dm::Context,
     defines: Option<dm::preprocessor::DefineHistory>,
     objtree: Arc<dm::objtree::ObjectTree>,
-    references_table: Option<find_references::ReferencesTable>,
+    references_table: background::Background<find_references::ReferencesTable>,
 
     annotations: HashMap<Url, (FileId, FileId, Rc<AnnotationTree>), RandomState>,
     diagnostics_set: HashSet<Url, RandomState>,
@@ -178,7 +179,7 @@ impl<'a> Engine<'a> {
             context,
             defines: None,
             objtree: Default::default(),
-            references_table: None,
+            references_table: Default::default(),
 
             annotations: Default::default(),
             diagnostics_set: Default::default(),
@@ -389,9 +390,13 @@ impl<'a> Engine<'a> {
             eprint!(" - object tree {}.{:03}s", elapsed.as_secs(), elapsed.subsec_millis());
         }
 
-        self.references_table = Some(find_references::ReferencesTable::new(&self.objtree));
-        let elapsed = start.elapsed(); start += elapsed;
-        eprint!(" - references {}.{:03}s", elapsed.as_secs(), elapsed.subsec_millis());
+        let references_objtree = self.objtree.clone();
+        self.references_table.spawn(move || {
+            let table = find_references::ReferencesTable::new(&references_objtree);
+            let elapsed = start.elapsed();
+            eprint!(" - references {}.{:03}s", elapsed.as_secs(), elapsed.subsec_millis());
+            table
+        });
 
         if ctx.config().langserver.dreamchecker && !fatal_errored {
             dreamchecker::run(&self.context, &self.objtree);
@@ -1535,7 +1540,8 @@ handle_method_call! {
 
         let mut result = &[][..];
         if let Some(id) = symbol_id {
-            if let Some(ref table) = self.references_table {
+            self.references_table.poll();
+            if let Some(table) = self.references_table.value() {
                 result = table.find_references(id, params.context.include_declaration);
             }
         }
@@ -1556,7 +1562,8 @@ handle_method_call! {
 
         let mut result = &[][..];
         if let Some(id) = symbol_id {
-            if let Some(ref table) = self.references_table {
+            self.references_table.poll();
+            if let Some(table) = self.references_table.value() {
                 result = table.find_implementations(id);
             }
         }