Fix language selection when saving new buffers as a single-file worktree

Co-Authored-By: Max Brunsfeld <maxbrunsfeld@gmail.com>

zed/src/editor/buffer.rs

@@ -719,11 +719,6 @@ impl Buffer {
         let text = self.visible_text.clone();
         let version = self.version.clone();
 
-        if let Some(language) = worktree.read(cx).languages().select_language(&path).cloned() {
-            self.language = Some(language);
-            self.reparse(cx);    
-        }
-
         let save_as = worktree.update(cx, |worktree, cx| {
             worktree
                 .as_local_mut()
@@ -734,6 +729,11 @@ impl Buffer {
         cx.spawn(|this, mut cx| async move {
             save_as.await.map(|new_file| {
                 this.update(&mut cx, |this, cx| {
+                    if let Some(language) = new_file.select_language(cx) {
+                        this.language = Some(language);
+                        this.reparse(cx);
+                    }
+
                     let mtime = new_file.mtime;
                     this.file = Some(new_file);
                     this.did_save(version, mtime, cx);

zed/src/language.rs

@@ -35,6 +35,10 @@ pub struct LanguageRegistry {
 }
 
 impl Language {
+    pub fn name(&self) -> &str {
+        self.config.name.as_str()
+    }
+
     pub fn highlight_map(&self) -> HighlightMap {
         self.highlight_map.lock().clone()
     }
@@ -133,27 +137,26 @@ mod tests {
 
         // matching file extension
         assert_eq!(
-            registry.select_language("zed/lib.rs").map(get_name),
+            registry.select_language("zed/lib.rs").map(|l| l.name()),
             Some("Rust")
         );
         assert_eq!(
-            registry.select_language("zed/lib.mk").map(get_name),
+            registry.select_language("zed/lib.mk").map(|l| l.name()),
             Some("Make")
         );
 
         // matching filename
         assert_eq!(
-            registry.select_language("zed/Makefile").map(get_name),
+            registry.select_language("zed/Makefile").map(|l| l.name()),
             Some("Make")
         );
 
         // matching suffix that is not the full file extension or filename
-        assert_eq!(registry.select_language("zed/cars").map(get_name), None);
+        assert_eq!(registry.select_language("zed/cars").map(|l| l.name()), None);
-        assert_eq!(registry.select_language("zed/a.cars").map(get_name), None);
+        assert_eq!(
-        assert_eq!(registry.select_language("zed/sumk").map(get_name), None);
+            registry.select_language("zed/a.cars").map(|l| l.name()),
-
+            None
-        fn get_name(language: &Arc<Language>) -> &str {
+        );
-            language.config.name.as_str()
+        assert_eq!(registry.select_language("zed/sumk").map(|l| l.name()), None);
-        }
     }
 }

zed/src/workspace.rs

@@ -1466,6 +1466,7 @@ mod tests {
         editor.update(&mut cx, |editor, cx| {
             assert!(!editor.is_dirty(cx.as_ref()));
             assert_eq!(editor.title(cx.as_ref()), "untitled");
+            assert!(editor.language(cx).is_none());
             editor.insert(&Insert("hi".into()), cx);
             assert!(editor.is_dirty(cx.as_ref()));
         });
@@ -1492,7 +1493,9 @@ mod tests {
             assert_eq!(editor.title(cx), "the-new-name.rs");
         });
         // The language is assigned based on the path
-        editor.read_with(&cx, |editor, cx| assert!(editor.language(cx).is_some()));
+        editor.read_with(&cx, |editor, cx| {
+            assert_eq!(editor.language(cx).unwrap().name(), "Rust")
+        });
 
         // Edit the file and save it again. This time, there is no filename prompt.
         editor.update(&mut cx, |editor, cx| {
@@ -1530,6 +1533,47 @@ mod tests {
         })
     }
 
+    #[gpui::test]
+    async fn test_setting_language_when_saving_as_single_file_worktree(
+        mut cx: gpui::TestAppContext,
+    ) {
+        let dir = TempDir::new("test-new-file").unwrap();
+        let app_state = cx.update(test_app_state);
+        let (_, workspace) = cx.add_window(|cx| Workspace::new(&app_state, cx));
+
+        // Create a new untitled buffer
+        let editor = workspace.update(&mut cx, |workspace, cx| {
+            workspace.open_new_file(&OpenNew(app_state.clone()), cx);
+            workspace
+                .active_item(cx)
+                .unwrap()
+                .to_any()
+                .downcast::<Editor>()
+                .unwrap()
+        });
+
+        editor.update(&mut cx, |editor, cx| {
+            assert!(editor.language(cx).is_none());
+            editor.insert(&Insert("hi".into()), cx);
+            assert!(editor.is_dirty(cx.as_ref()));
+        });
+
+        // Save the buffer. This prompts for a filename.
+        workspace.update(&mut cx, |workspace, cx| {
+            workspace.save_active_item(&Save, cx)
+        });
+        cx.simulate_new_path_selection(|_| Some(dir.path().join("the-new-name.rs")));
+
+        editor
+            .condition(&cx, |editor, cx| !editor.is_dirty(cx))
+            .await;
+
+        // The language is assigned based on the path
+        editor.read_with(&cx, |editor, cx| {
+            assert_eq!(editor.language(cx).unwrap().name(), "Rust")
+        });
+    }
+
     #[gpui::test]
     async fn test_new_empty_workspace(mut cx: gpui::TestAppContext) {
         cx.update(init);

zed/src/worktree.rs

@@ -6,7 +6,7 @@ use crate::{
     fs::{self, Fs},
     fuzzy,
     fuzzy::CharBag,
-    language::LanguageRegistry,
+    language::{Language, LanguageRegistry},
     rpc::{self, proto},
     time::{self, ReplicaId},
     util::{Bias, TryFutureExt},
@@ -268,13 +268,6 @@ impl Worktree {
         }
     }
 
-    pub fn languages(&self) -> &Arc<LanguageRegistry> {
-        match self {
-            Worktree::Local(worktree) => &worktree.languages,
-            Worktree::Remote(worktree) => &worktree.languages,
-        }
-    }
-
     pub fn snapshot(&self) -> Snapshot {
         match self {
             Worktree::Local(worktree) => worktree.snapshot(),
@@ -784,7 +777,6 @@ impl LocalWorktree {
             }
         });
 
-        let languages = self.languages.clone();
         let path = Arc::from(path);
         cx.spawn(|this, mut cx| async move {
             if let Some(existing_buffer) = existing_buffer {
@@ -793,8 +785,8 @@ impl LocalWorktree {
                 let (file, contents) = this
                     .update(&mut cx, |this, cx| this.as_local().unwrap().load(&path, cx))
                     .await?;
-                let language = languages.select_language(&path).cloned();
                 let buffer = cx.add_model(|cx| {
+                    let language = file.select_language(cx);
                     Buffer::from_history(0, History::new(contents.into()), Some(file), language, cx)
                 });
                 this.update(&mut cx, |this, _| {
@@ -1127,7 +1119,6 @@ impl RemoteWorktree {
         });
 
         let rpc = self.rpc.clone();
-        let languages = self.languages.clone();
         let replica_id = self.replica_id;
         let remote_worktree_id = self.remote_id;
         let path = path.to_string_lossy().to_string();
@@ -1139,7 +1130,7 @@ impl RemoteWorktree {
                     .read_with(&cx, |tree, _| tree.entry_for_path(&path).cloned())
                     .ok_or_else(|| anyhow!("file does not exist"))?;
                 let file = File::new(entry.id, handle, entry.path, entry.mtime);
-                let language = languages.select_language(&path).cloned();
+                let language = cx.read(|cx| file.select_language(cx));
                 let response = rpc
                     .request(proto::OpenBuffer {
                         worktree_id: remote_worktree_id as u64,
@@ -1616,6 +1607,18 @@ impl File {
         self.worktree.read(cx).abs_path.join(&self.path)
     }
 
+    pub fn select_language(&self, cx: &AppContext) -> Option<Arc<Language>> {
+        let worktree = self.worktree.read(cx);
+        let mut full_path = PathBuf::new();
+        full_path.push(worktree.root_name());
+        full_path.push(&self.path);
+        let languages = match self.worktree.read(cx) {
+            Worktree::Local(worktree) => &worktree.languages,
+            Worktree::Remote(worktree) => &worktree.languages,
+        };
+        languages.select_language(&full_path).cloned()
+    }
+
     /// Returns the last component of this handle's absolute path. If this handle refers to the root
     /// of its worktree, then this method will return the name of the worktree itself.
     pub fn file_name<'a>(&'a self, cx: &'a AppContext) -> Option<OsString> {