Add HOT_RELOAD_BUILTIN_TREE_SITTER_QUERIES env var for dev

crates/languages/README.md

@@ -0,0 +1,11 @@
+This crate contains Zed's builtin languages. Eventually these will be migrated to the extension system.
+
+When updating the `*.scm` TreeSitter query files, the following can be helpful:
+
+```sh
+HOT_RELOAD_BUILTIN_TREE_SITTER_QUERIES=. cargo run
+```
+
+The implementation of this is fairly inefficient - on every file change it reloads all builtin languages and resets the language on all buffers.
+
+This also works with release builds of Zed, and when run from some other directory it can be set to the path of the Zed repository that has the queries.

crates/languages/src/lib.rs

@@ -394,7 +394,30 @@ fn load_config(name: &str) -> LanguageConfig {
     config
 }
 
+/// Name of the environment variable to provide in order to hot-reload the TreeSitter queries. When
+/// running from the Zed repository root, this should just be set to ".". Otherwise, the Zed
+/// repository root can be specified.
+pub const HOT_RELOAD_BUILTIN_TREE_SITTER_QUERIES: &str = "HOT_RELOAD_BUILTIN_TREE_SITTER_QUERIES";
+pub const TREE_SITTER_QUERIES_SOURCE_PATH: &str = "crates/languages/src/";
+
 fn load_queries(name: &str) -> LanguageQueries {
+    if let Ok(zed_repo_path) = std::env::var(HOT_RELOAD_BUILTIN_TREE_SITTER_QUERIES) {
+        let zed_repo_path = std::path::PathBuf::from(zed_repo_path);
+        match load_queries_from_fs(name, zed_repo_path) {
+            Ok(queries) => return queries,
+            Err(err) => {
+                log::error!(
+                    "Failed to load queries from files (due to {}): {}",
+                    HOT_RELOAD_BUILTIN_TREE_SITTER_QUERIES,
+                    err
+                );
+            }
+        }
+    }
+    load_queries_from_assets(name)
+}
+
+fn load_queries_from_assets(name: &str) -> LanguageQueries {
     let mut result = LanguageQueries::default();
     for path in LanguageDir::iter() {
         if let Some(remainder) = path.strip_prefix(name).and_then(|p| p.strip_prefix('/')) {
@@ -414,3 +437,37 @@ fn load_queries(name: &str) -> LanguageQueries {
     }
     result
 }
+
+fn load_queries_from_fs(
+    name: &str,
+    zed_repo_path: std::path::PathBuf,
+) -> anyhow::Result<LanguageQueries> {
+    let mut queries_dir = zed_repo_path;
+    queries_dir.push(TREE_SITTER_QUERIES_SOURCE_PATH);
+    queries_dir.push(name);
+    if !std::fs::exists(&queries_dir)? {
+        anyhow::bail!("{} does not exist", queries_dir.display());
+    }
+    let mut result = LanguageQueries::default();
+    for entry in std::fs::read_dir(queries_dir)? {
+        let entry = entry?;
+        let path = entry.path();
+        if path.extension().and_then(|ext| ext.to_str()) != Some("scm") {
+            continue;
+        }
+        for (name, query) in QUERY_FILENAME_PREFIXES {
+            if path
+                .file_name()
+                .and_then(|file_name| file_name.to_str())
+                .is_some_and(|file_name| file_name.starts_with(name))
+            {
+                let contents = std::fs::read_to_string(&path)?;
+                match query(&mut result) {
+                    None => *query(&mut result) = Some(contents.into()),
+                    Some(r) => r.to_mut().push_str(contents.as_ref()),
+                }
+            }
+        }
+    }
+    Ok(result)
+}

crates/zed/src/main.rs

@@ -468,6 +468,11 @@ pub fn main() {
         debug_adapter_extension::init(extension_host_proxy.clone(), cx);
         language::init(cx);
         languages::init(languages.clone(), node_runtime.clone(), cx);
+        zed::hot_reload_builtin_treesitter_queries::init(
+            languages.clone(),
+            node_runtime.clone(),
+            cx,
+        );
         let user_store = cx.new(|cx| UserStore::new(client.clone(), cx));
         let workspace_store = cx.new(|cx| WorkspaceStore::new(client.clone(), cx));
 

crates/zed/src/zed.rs

@@ -1,6 +1,7 @@
 mod app_menus;
 pub mod component_preview;
 pub mod edit_prediction_registry;
+pub(crate) mod hot_reload_builtin_treesitter_queries;
 #[cfg(target_os = "macos")]
 pub(crate) mod mac_only_instance;
 mod migrate;

crates/zed/src/zed/hot_reload_builtin_treesitter_queries.rs

@@ -0,0 +1,122 @@
+use anyhow::Result;
+use fs::Fs;
+use futures::{StreamExt, channel::mpsc};
+use gpui::{App, AppContext, AsyncApp, Entity};
+use language::{Buffer, LanguageRegistry};
+use languages::{HOT_RELOAD_BUILTIN_TREE_SITTER_QUERIES, TREE_SITTER_QUERIES_SOURCE_PATH};
+use node_runtime::NodeRuntime;
+use std::{env, path::PathBuf, sync::Arc, time::Duration};
+use util::ResultExt as _;
+use workspace::Workspace;
+
+pub(crate) fn init(languages: Arc<LanguageRegistry>, node: NodeRuntime, cx: &mut App) {
+    let Ok(zed_repo_path) = env::var(HOT_RELOAD_BUILTIN_TREE_SITTER_QUERIES) else {
+        return;
+    };
+    let zed_repo_path = PathBuf::from(zed_repo_path);
+    let fs = <dyn Fs>::global(cx);
+    let (tx, mut rx) = mpsc::unbounded();
+    cx.spawn(
+        async move |cx| match setup_watches(zed_repo_path, tx, fs, cx).await {
+            Ok(()) => {
+                while let Some(()) = rx.next().await {
+                    cx.update(|cx| reload_treesitter_queries(languages.clone(), node.clone(), cx))
+                        .ok();
+                }
+            }
+            Err(err) => log::error!(
+                "Failed to setup watches to handle {}: {}",
+                err,
+                HOT_RELOAD_BUILTIN_TREE_SITTER_QUERIES
+            ),
+        },
+    )
+    .detach();
+}
+
+async fn setup_watches(
+    zed_repo_path: PathBuf,
+    tx: mpsc::UnboundedSender<()>,
+    fs: Arc<dyn Fs>,
+    cx: &mut AsyncApp,
+) -> Result<()> {
+    let mut queries_dir = zed_repo_path;
+    queries_dir.push(TREE_SITTER_QUERIES_SOURCE_PATH);
+    let mut paths = fs.read_dir(&queries_dir).await?;
+    while let Some(path) = paths.next().await {
+        let path = path?;
+        if fs.is_dir(&path).await {
+            let tx = tx.clone();
+            let fs = fs.clone();
+            cx.background_spawn(async move {
+                let (events, _) = fs.watch(&path, Duration::from_millis(100)).await;
+                futures::pin_mut!(events);
+                while let Some(event_batch) = events.next().await {
+                    let query_modified = event_batch
+                        .into_iter()
+                        .any(|event| event.path.extension().is_some_and(|ext| ext == "scm"));
+                    if query_modified {
+                        tx.unbounded_send(()).log_err();
+                    }
+                }
+            })
+            .detach();
+        }
+    }
+    log::info!(
+        "Watching {} for changes due to use of {}",
+        queries_dir.display(),
+        HOT_RELOAD_BUILTIN_TREE_SITTER_QUERIES
+    );
+    Ok(())
+}
+
+fn reload_treesitter_queries(languages: Arc<LanguageRegistry>, node: NodeRuntime, cx: &mut App) {
+    languages::init(languages.clone(), node, cx);
+    for window in cx.windows() {
+        if let Some(workspace) = window.downcast::<Workspace>()
+            && let Ok(workspace) = workspace.entity(cx)
+        {
+            let buffers = workspace
+                .read(cx)
+                .project()
+                .read(cx)
+                .buffer_store()
+                .read(cx)
+                .buffers()
+                .collect::<Vec<_>>();
+            for buffer in buffers {
+                reset_buffer_language(buffer, languages.clone(), cx);
+            }
+        }
+    }
+}
+
+fn reset_buffer_language(buffer: Entity<Buffer>, languages: Arc<LanguageRegistry>, cx: &App) {
+    if let Some(old_language) = buffer.read(cx).language() {
+        let old_language_name = old_language.name().to_string();
+        let new_language = cx.background_spawn({
+            let old_language_name = old_language_name.clone();
+            async move { languages.language_for_name(&old_language_name).await }
+        });
+        cx.spawn(async move |cx_wat| match new_language.await {
+            Ok(new_language) => {
+                buffer
+                    .update(cx_wat, |buffer, cx| {
+                        buffer.set_language(Some(new_language), cx);
+                    })
+                    .ok();
+            }
+            Err(err) => {
+                log::warn!(
+                    "While handling {}, failed to update buffer with language \"{}\" \
+                    (probably just not a built-in language): {}",
+                    HOT_RELOAD_BUILTIN_TREE_SITTER_QUERIES,
+                    old_language_name,
+                    err
+                );
+            }
+        })
+        .detach();
+    }
+}