Live-reload tree-sitter queries in development

crates/language/src/language.rs

@@ -34,7 +34,7 @@ use std::{
     fmt::Debug,
     hash::Hash,
     mem,
-    ops::Range,
+    ops::{Not, Range},
     path::{Path, PathBuf},
     str,
     sync::{
@@ -500,6 +500,7 @@ struct AvailableLanguage {
     grammar: tree_sitter::Language,
     lsp_adapters: Vec<Arc<dyn LspAdapter>>,
     get_queries: fn(&str) -> LanguageQueries,
+    loaded: bool,
 }
 
 pub struct LanguageRegistry {
@@ -527,6 +528,7 @@ struct LanguageRegistryState {
     subscription: (watch::Sender<()>, watch::Receiver<()>),
     theme: Option<Arc<Theme>>,
     version: usize,
+    reload_count: usize,
 }
 
 pub struct PendingLanguageServer {
@@ -547,6 +549,7 @@ impl LanguageRegistry {
                 subscription: watch::channel(),
                 theme: Default::default(),
                 version: 0,
+                reload_count: 0,
             }),
             language_server_download_dir: None,
             lsp_binary_statuses_tx,
@@ -566,6 +569,14 @@ impl LanguageRegistry {
         self.executor = Some(executor);
     }
 
+    /// Clear out all of the loaded languages and reload them from scratch.
+    ///
+    /// This is useful in development, when queries have changed.
+    #[cfg(debug_assertions)]
+    pub fn reload(&self) {
+        self.state.write().reload();
+    }
+
     pub fn register(
         &self,
         path: &'static str,
@@ -582,6 +593,7 @@ impl LanguageRegistry {
             grammar,
             lsp_adapters,
             get_queries,
+            loaded: false,
         });
     }
 
@@ -590,7 +602,7 @@ impl LanguageRegistry {
         let mut result = state
             .available_languages
             .iter()
-            .map(|l| l.config.name.to_string())
+            .filter_map(|l| l.loaded.not().then_some(l.config.name.to_string()))
             .chain(state.languages.iter().map(|l| l.config.name.to_string()))
             .collect::<Vec<_>>();
         result.sort_unstable_by_key(|language_name| language_name.to_lowercase());
@@ -603,6 +615,7 @@ impl LanguageRegistry {
             state
                 .available_languages
                 .iter()
+                .filter(|l| !l.loaded)
                 .flat_map(|l| l.lsp_adapters.clone())
                 .chain(
                     state
@@ -639,10 +652,17 @@ impl LanguageRegistry {
         self.state.read().subscription.1.clone()
     }
 
+    /// The number of times that the registry has been changed,
+    /// by adding languages or reloading.
     pub fn version(&self) -> usize {
         self.state.read().version
     }
 
+    /// The number of times that the registry has been reloaded.
+    pub fn reload_count(&self) -> usize {
+        self.state.read().reload_count
+    }
+
     pub fn set_theme(&self, theme: Arc<Theme>) {
         let mut state = self.state.write();
         state.theme = Some(theme.clone());
@@ -721,7 +741,7 @@ impl LanguageRegistry {
             if let Some(language) = state
                 .available_languages
                 .iter()
-                .find(|l| callback(&l.config))
+                .find(|l| !l.loaded && callback(&l.config))
                 .cloned()
             {
                 let txs = state
@@ -743,9 +763,7 @@ impl LanguageRegistry {
                                         let language = Arc::new(language);
                                         let mut state = this.state.write();
                                         state.add(language.clone());
-                                        state
-                                            .available_languages
-                                            .retain(|language| language.id != id);
+                                        state.mark_language_loaded(id);
                                         if let Some(mut txs) = state.loading_languages.remove(&id) {
                                             for tx in txs.drain(..) {
                                                 let _ = tx.send(Ok(language.clone()));
@@ -754,9 +772,7 @@ impl LanguageRegistry {
                                     }
                                     Err(err) => {
                                         let mut state = this.state.write();
-                                        state
-                                            .available_languages
-                                            .retain(|language| language.id != id);
+                                        state.mark_language_loaded(id);
                                         if let Some(mut txs) = state.loading_languages.remove(&id) {
                                             for tx in txs.drain(..) {
                                                 let _ = tx.send(Err(anyhow!(
@@ -905,6 +921,28 @@ impl LanguageRegistryState {
         self.version += 1;
         *self.subscription.0.borrow_mut() = ();
     }
+
+    #[cfg(debug_assertions)]
+    fn reload(&mut self) {
+        self.languages.clear();
+        self.version += 1;
+        self.reload_count += 1;
+        for language in &mut self.available_languages {
+            language.loaded = false;
+        }
+        *self.subscription.0.borrow_mut() = ();
+    }
+
+    /// Mark the given language a having been loaded, so that the
+    /// language registry won't try to load it again.
+    fn mark_language_loaded(&mut self, id: AvailableLanguageId) {
+        for language in &mut self.available_languages {
+            if language.id == id {
+                language.loaded = true;
+                break;
+            }
+        }
+    }
 }
 
 #[cfg(any(test, feature = "test-support"))]

crates/project/src/project.rs

@@ -523,7 +523,7 @@ impl Project {
                 _subscriptions: vec![
                     cx.observe_global::<SettingsStore, _>(Self::on_settings_changed)
                 ],
-                _maintain_buffer_languages: Self::maintain_buffer_languages(&languages, cx),
+                _maintain_buffer_languages: Self::maintain_buffer_languages(languages.clone(), cx),
                 _maintain_workspace_config: Self::maintain_workspace_config(languages.clone(), cx),
                 active_entry: None,
                 languages,
@@ -592,7 +592,7 @@ impl Project {
                 active_entry: None,
                 collaborators: Default::default(),
                 join_project_response_message_id: response.message_id,
-                _maintain_buffer_languages: Self::maintain_buffer_languages(&languages, cx),
+                _maintain_buffer_languages: Self::maintain_buffer_languages(languages.clone(), cx),
                 _maintain_workspace_config: Self::maintain_workspace_config(languages.clone(), cx),
                 languages,
                 user_store: user_store.clone(),
@@ -2238,13 +2238,34 @@ impl Project {
     }
 
     fn maintain_buffer_languages(
-        languages: &LanguageRegistry,
+        languages: Arc<LanguageRegistry>,
         cx: &mut ModelContext<Project>,
     ) -> Task<()> {
         let mut subscription = languages.subscribe();
+        let mut prev_reload_count = languages.reload_count();
         cx.spawn_weak(|project, mut cx| async move {
             while let Some(()) = subscription.next().await {
                 if let Some(project) = project.upgrade(&cx) {
+                    // If the language registry has been reloaded, then remove and
+                    // re-assign the languages on all open buffers.
+                    let reload_count = languages.reload_count();
+                    if reload_count > prev_reload_count {
+                        prev_reload_count = reload_count;
+                        project.update(&mut cx, |this, cx| {
+                            let buffers = this
+                                .opened_buffers
+                                .values()
+                                .filter_map(|b| b.upgrade(cx))
+                                .collect::<Vec<_>>();
+                            for buffer in buffers {
+                                if let Some(f) = File::from_dyn(buffer.read(cx).file()).cloned() {
+                                    this.unregister_buffer_from_language_servers(&buffer, &f, cx);
+                                    buffer.update(cx, |buffer, cx| buffer.set_language(None, cx));
+                                }
+                            }
+                        });
+                    }
+
                     project.update(&mut cx, |project, cx| {
                         let mut plain_text_buffers = Vec::new();
                         let mut buffers_with_unknown_injections = Vec::new();

crates/zed/src/main.rs

@@ -160,6 +160,8 @@ fn main() {
         ai::init(cx);
 
         cx.spawn(|cx| watch_themes(fs.clone(), cx)).detach();
+        cx.spawn(|_| watch_languages(fs.clone(), languages.clone()))
+            .detach();
 
         languages.set_theme(theme::current(cx).clone());
         cx.observe_global::<SettingsStore, _>({
@@ -660,11 +662,30 @@ async fn watch_themes(fs: Arc<dyn Fs>, mut cx: AsyncAppContext) -> Option<()> {
     Some(())
 }
 
+#[cfg(debug_assertions)]
+async fn watch_languages(fs: Arc<dyn Fs>, languages: Arc<LanguageRegistry>) -> Option<()> {
+    let mut events = fs
+        .watch(
+            "crates/zed/src/languages".as_ref(),
+            Duration::from_millis(100),
+        )
+        .await;
+    while (events.next().await).is_some() {
+        languages.reload();
+    }
+    Some(())
+}
+
 #[cfg(not(debug_assertions))]
 async fn watch_themes(_fs: Arc<dyn Fs>, _cx: AsyncAppContext) -> Option<()> {
     None
 }
 
+#[cfg(not(debug_assertions))]
+async fn watch_languages(_: Arc<dyn Fs>, _: Arc<LanguageRegistry>) -> Option<()> {
+    None
+}
+
 fn connect_to_cli(
     server_name: &str,
 ) -> Result<(mpsc::Receiver<CliRequest>, IpcSender<CliResponse>)> {