Unload languages when uninstalling their extension (#7743)

Release Notes:

- N/A

Co-authored-by: Marshall <marshall@zed.dev>

crates/extension/src/extension_store.rs

@@ -311,6 +311,13 @@ impl ExtensionStore {
     }
 
     fn manifest_updated(&mut self, manifest: Manifest, cx: &mut ModelContext<Self>) {
+        let old_manifest = self.manifest.read();
+        let language_names = old_manifest.languages.keys().cloned().collect::<Vec<_>>();
+        let grammar_names = old_manifest.grammars.keys().cloned().collect::<Vec<_>>();
+        drop(old_manifest);
+        self.language_registry
+            .remove_languages(&language_names, &grammar_names);
+
         self.language_registry
             .register_wasm_grammars(manifest.grammars.iter().map(|(grammar_name, grammar)| {
                 let mut grammar_path = self.extensions_dir.clone();

crates/extension/src/extension_store_test.rs

@@ -302,6 +302,10 @@ async fn test_extension_store(cx: &mut TestAppContext) {
             language_registry.language_names(),
             ["ERB", "Plain Text", "Ruby"]
         );
+        assert_eq!(
+            language_registry.grammar_names(),
+            ["embedded_template".into(), "ruby".into()]
+        );
         assert_eq!(
             theme_registry.list_names(false),
             [
@@ -319,4 +323,27 @@ async fn test_extension_store(cx: &mut TestAppContext) {
         assert_eq!(fs.read_dir_call_count(), prev_fs_read_dir_call_count);
         assert_eq!(fs.metadata_call_count(), prev_fs_metadata_call_count + 2);
     });
+
+    store
+        .update(cx, |store, cx| {
+            store.uninstall_extension("zed-ruby".into(), cx)
+        })
+        .await
+        .unwrap();
+
+    expected_manifest.extensions.remove("zed-ruby");
+    expected_manifest.languages.remove("Ruby");
+    expected_manifest.languages.remove("ERB");
+    expected_manifest.grammars.remove("ruby");
+    expected_manifest.grammars.remove("embedded_template");
+
+    store.read_with(cx, |store, _| {
+        let manifest = store.manifest.read();
+        assert_eq!(manifest.grammars, expected_manifest.grammars);
+        assert_eq!(manifest.languages, expected_manifest.languages);
+        assert_eq!(manifest.themes, expected_manifest.themes);
+
+        assert_eq!(language_registry.language_names(), ["Plain Text"]);
+        assert_eq!(language_registry.grammar_names(), []);
+    });
 }

crates/language/src/language.rs

@@ -83,8 +83,8 @@ thread_local! {
 }
 
 lazy_static! {
+    static ref NEXT_LANGUAGE_ID: AtomicUsize = Default::default();
     static ref NEXT_GRAMMAR_ID: AtomicUsize = Default::default();
-
     static ref WASM_ENGINE: wasmtime::Engine = wasmtime::Engine::default();
 
     /// A shared grammar for plain text, exposed for reuse by downstream crates.
@@ -613,7 +613,17 @@ pub struct BracketPair {
     pub newline: bool,
 }
 
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy)]
+pub(crate) struct LanguageId(usize);
+
+impl LanguageId {
+    pub(crate) fn new() -> Self {
+        Self(NEXT_LANGUAGE_ID.fetch_add(1, SeqCst))
+    }
+}
+
 pub struct Language {
+    pub(crate) id: LanguageId,
     pub(crate) config: LanguageConfig,
     pub(crate) grammar: Option<Arc<Grammar>>,
     pub(crate) adapters: Vec<Arc<CachedLspAdapter>>,
@@ -625,8 +635,17 @@ pub struct Language {
     )>,
 }
 
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy)]
+pub struct GrammarId(pub usize);
+
+impl GrammarId {
+    pub(crate) fn new() -> Self {
+        Self(NEXT_GRAMMAR_ID.fetch_add(1, SeqCst))
+    }
+}
+
 pub struct Grammar {
-    id: usize,
+    id: GrammarId,
     pub ts_language: tree_sitter::Language,
     pub(crate) error_query: Query,
     pub(crate) highlights_query: Option<Query>,
@@ -697,11 +716,24 @@ struct BracketConfig {
 
 impl Language {
     pub fn new(config: LanguageConfig, ts_language: Option<tree_sitter::Language>) -> Self {
+        Self::new_with_id(
+            LanguageId(NEXT_LANGUAGE_ID.fetch_add(1, SeqCst)),
+            config,
+            ts_language,
+        )
+    }
+
+    fn new_with_id(
+        id: LanguageId,
+        config: LanguageConfig,
+        ts_language: Option<tree_sitter::Language>,
+    ) -> Self {
         Self {
+            id,
             config,
             grammar: ts_language.map(|ts_language| {
                 Arc::new(Grammar {
-                    id: NEXT_GRAMMAR_ID.fetch_add(1, SeqCst),
+                    id: GrammarId::new(),
                     highlights_query: None,
                     brackets_config: None,
                     outline_config: None,
@@ -726,10 +758,6 @@ impl Language {
         &self.adapters
     }
 
-    pub fn id(&self) -> Option<usize> {
-        self.grammar.as_ref().map(|g| g.id)
-    }
-
     pub fn with_queries(mut self, queries: LanguageQueries) -> Result<Self> {
         if let Some(query) = queries.highlights {
             self = self
@@ -1237,13 +1265,13 @@ impl LanguageScope {
 
 impl Hash for Language {
     fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
-        self.id().hash(state)
+        self.id.hash(state)
     }
 }
 
 impl PartialEq for Language {
     fn eq(&self, other: &Self) -> bool {
-        self.id().eq(&other.id())
+        self.id.eq(&other.id)
     }
 }
 
@@ -1258,7 +1286,7 @@ impl Debug for Language {
 }
 
 impl Grammar {
-    pub fn id(&self) -> usize {
+    pub fn id(&self) -> GrammarId {
         self.id
     }
 

crates/language/src/language_registry.rs

@@ -1,6 +1,6 @@
 use crate::{
-    CachedLspAdapter, Language, LanguageConfig, LanguageMatcher, LanguageServerName, LspAdapter,
+    CachedLspAdapter, Language, LanguageConfig, LanguageId, LanguageMatcher, LanguageServerName,
-    LspAdapterDelegate, PARSER, PLAIN_TEXT,
+    LspAdapter, LspAdapterDelegate, PARSER, PLAIN_TEXT,
 };
 use anyhow::{anyhow, Context as _, Result};
 use collections::{hash_map, HashMap};
@@ -43,8 +43,7 @@ struct LanguageRegistryState {
     languages: Vec<Arc<Language>>,
     available_languages: Vec<AvailableLanguage>,
     grammars: HashMap<Arc<str>, AvailableGrammar>,
-    next_available_language_id: AvailableLanguageId,
+    loading_languages: HashMap<LanguageId, Vec<oneshot::Sender<Result<Arc<Language>>>>>,
-    loading_languages: HashMap<AvailableLanguageId, Vec<oneshot::Sender<Result<Arc<Language>>>>>,
     subscription: (watch::Sender<()>, watch::Receiver<()>),
     theme: Option<Arc<Theme>>,
     version: usize,
@@ -68,7 +67,7 @@ pub struct PendingLanguageServer {
 
 #[derive(Clone)]
 struct AvailableLanguage {
-    id: AvailableLanguageId,
+    id: LanguageId,
     name: Arc<str>,
     grammar: Option<Arc<str>>,
     matcher: LanguageMatcher,
@@ -77,8 +76,6 @@ struct AvailableLanguage {
     loaded: bool,
 }
 
-type AvailableLanguageId = usize;
-
 enum AvailableGrammar {
     Native(tree_sitter::Language),
     Loaded(PathBuf, tree_sitter::Language),
@@ -126,7 +123,6 @@ impl LanguageRegistry {
                 languages: vec![PLAIN_TEXT.clone()],
                 available_languages: Default::default(),
                 grammars: Default::default(),
-                next_available_language_id: 0,
                 loading_languages: Default::default(),
                 subscription: watch::channel(),
                 theme: Default::default(),
@@ -160,6 +156,17 @@ impl LanguageRegistry {
         self.state.write().reload_languages(languages, grammars);
     }
 
+    /// Removes the specified languages and grammars from the registry.
+    pub fn remove_languages(
+        &self,
+        languages_to_remove: &[Arc<str>],
+        grammars_to_remove: &[Arc<str>],
+    ) {
+        self.state
+            .write()
+            .remove_languages(languages_to_remove, grammars_to_remove)
+    }
+
     #[cfg(any(feature = "test-support", test))]
     pub fn register_test_language(&self, config: LanguageConfig) {
         self.register_language(
@@ -194,7 +201,7 @@ impl LanguageRegistry {
         }
 
         state.available_languages.push(AvailableLanguage {
-            id: post_inc(&mut state.next_available_language_id),
+            id: LanguageId::new(),
             name,
             grammar: grammar_name,
             matcher,
@@ -241,6 +248,13 @@ impl LanguageRegistry {
         result
     }
 
+    pub fn grammar_names(&self) -> Vec<Arc<str>> {
+        let state = self.state.read();
+        let mut result = state.grammars.keys().cloned().collect::<Vec<_>>();
+        result.sort_unstable_by_key(|grammar_name| grammar_name.to_lowercase());
+        result
+    }
+
     pub fn add(&self, language: Arc<Language>) {
         self.state.write().add(language);
     }
@@ -358,7 +372,7 @@ impl LanguageRegistry {
                                         None
                                     };
 
-                                    Language::new(config, grammar)
+                                    Language::new_with_id(id, config, grammar)
                                         .with_lsp_adapters(language.lsp_adapters)
                                         .await
                                         .with_queries(queries)
@@ -660,6 +674,22 @@ impl LanguageRegistryState {
         *self.subscription.0.borrow_mut() = ();
     }
 
+    fn remove_languages(
+        &mut self,
+        languages_to_remove: &[Arc<str>],
+        grammars_to_remove: &[Arc<str>],
+    ) {
+        self.languages
+            .retain(|language| !languages_to_remove.contains(&language.name()));
+        self.available_languages
+            .retain(|language| !languages_to_remove.contains(&language.name));
+        self.grammars
+            .retain(|name, _| !grammars_to_remove.contains(&name));
+        self.version += 1;
+        self.reload_count += 1;
+        *self.subscription.0.borrow_mut() = ();
+    }
+
     fn reload_languages(
         &mut self,
         languages_to_reload: &[Arc<str>],
@@ -701,7 +731,7 @@ impl LanguageRegistryState {
 
     /// Mark the given language as having been loaded, so that the
     /// language registry won't try to load it again.
-    fn mark_language_loaded(&mut self, id: AvailableLanguageId) {
+    fn mark_language_loaded(&mut self, id: LanguageId) {
         for language in &mut self.available_languages {
             if language.id == id {
                 language.loaded = true;

crates/language/src/syntax_map.rs

@@ -1,7 +1,7 @@
 #[cfg(test)]
 mod syntax_map_tests;
 
-use crate::{Grammar, InjectionConfig, Language, LanguageRegistry};
+use crate::{Grammar, InjectionConfig, Language, LanguageId, LanguageRegistry};
 use collections::HashMap;
 use futures::FutureExt;
 use parking_lot::Mutex;
@@ -102,9 +102,9 @@ enum SyntaxLayerContent {
 }
 
 impl SyntaxLayerContent {
-    fn language_id(&self) -> Option<usize> {
+    fn language_id(&self) -> Option<LanguageId> {
         match self {
-            SyntaxLayerContent::Parsed { language, .. } => language.id(),
+            SyntaxLayerContent::Parsed { language, .. } => Some(language.id),
             SyntaxLayerContent::Pending { .. } => None,
         }
     }
@@ -144,7 +144,7 @@ struct SyntaxLayerSummary {
     max_depth: usize,
     range: Range<Anchor>,
     last_layer_range: Range<Anchor>,
-    last_layer_language: Option<usize>,
+    last_layer_language: Option<LanguageId>,
     contains_unknown_injections: bool,
 }
 
@@ -152,7 +152,7 @@ struct SyntaxLayerSummary {
 struct SyntaxLayerPosition {
     depth: usize,
     range: Range<Anchor>,
-    language: Option<usize>,
+    language: Option<LanguageId>,
 }
 
 #[derive(Clone, Debug)]
@@ -182,9 +182,9 @@ enum ParseStepLanguage {
 }
 
 impl ParseStepLanguage {
-    fn id(&self) -> Option<usize> {
+    fn id(&self) -> Option<LanguageId> {
         match self {
-            ParseStepLanguage::Loaded { language } => language.id(),
+            ParseStepLanguage::Loaded { language } => Some(language.id),
             ParseStepLanguage::Pending { .. } => None,
         }
     }