Fix panic loading language queries (#9506)

Conrad Irwin created 2 years ago

Release Notes:

- Fixed a panic when loading a language with queries but no grammar

Change summary

crates/language/src/language.rs          | 47 +++++++++++++++++++------
crates/language/src/language_registry.rs | 17 ++++----
2 files changed, 45 insertions(+), 19 deletions(-)

Detailed changes

crates/language/src/language.rs

@@ -940,13 +940,17 @@ impl Language {
     }
 
     pub fn with_highlights_query(mut self, source: &str) -> Result<Self> {
-        let grammar = self.grammar_mut();
+        let grammar = self
+            .grammar_mut()
+            .ok_or_else(|| anyhow!("cannot mutate grammar"))?;
         grammar.highlights_query = Some(Query::new(&grammar.ts_language, source)?);
         Ok(self)
     }
 
     pub fn with_outline_query(mut self, source: &str) -> Result<Self> {
-        let grammar = self.grammar_mut();
+        let grammar = self
+            .grammar_mut()
+            .ok_or_else(|| anyhow!("cannot mutate grammar"))?;
         let query = Query::new(&grammar.ts_language, source)?;
         let mut item_capture_ix = None;
         let mut name_capture_ix = None;
@@ -974,7 +978,9 @@ impl Language {
     }
 
     pub fn with_embedding_query(mut self, source: &str) -> Result<Self> {
-        let grammar = self.grammar_mut();
+        let grammar = self
+            .grammar_mut()
+            .ok_or_else(|| anyhow!("cannot mutate grammar"))?;
         let query = Query::new(&grammar.ts_language, source)?;
         let mut item_capture_ix = None;
         let mut name_capture_ix = None;
@@ -1005,7 +1011,9 @@ impl Language {
     }
 
     pub fn with_brackets_query(mut self, source: &str) -> Result<Self> {
-        let grammar = self.grammar_mut();
+        let grammar = self
+            .grammar_mut()
+            .ok_or_else(|| anyhow!("cannot mutate grammar"))?;
         let query = Query::new(&grammar.ts_language, source)?;
         let mut open_capture_ix = None;
         let mut close_capture_ix = None;
@@ -1027,7 +1035,9 @@ impl Language {
     }
 
     pub fn with_indents_query(mut self, source: &str) -> Result<Self> {
-        let grammar = self.grammar_mut();
+        let grammar = self
+            .grammar_mut()
+            .ok_or_else(|| anyhow!("cannot mutate grammar"))?;
         let query = Query::new(&grammar.ts_language, source)?;
         let mut indent_capture_ix = None;
         let mut start_capture_ix = None;
@@ -1055,7 +1065,9 @@ impl Language {
     }
 
     pub fn with_injection_query(mut self, source: &str) -> Result<Self> {
-        let grammar = self.grammar_mut();
+        let grammar = self
+            .grammar_mut()
+            .ok_or_else(|| anyhow!("cannot mutate grammar"))?;
         let query = Query::new(&grammar.ts_language, source)?;
         let mut language_capture_ix = None;
         let mut content_capture_ix = None;
@@ -1095,7 +1107,13 @@ impl Language {
     }
 
     pub fn with_override_query(mut self, source: &str) -> anyhow::Result<Self> {
-        let query = Query::new(&self.grammar_mut().ts_language, source)?;
+        let query = {
+            let grammar = self
+                .grammar
+                .as_ref()
+                .ok_or_else(|| anyhow!("no grammar for language"))?;
+            Query::new(&grammar.ts_language, source)?
+        };
 
         let mut override_configs_by_id = HashMap::default();
         for (ix, name) in query.capture_names().iter().enumerate() {
@@ -1159,7 +1177,11 @@ impl Language {
         }
 
         self.config.brackets.disabled_scopes_by_bracket_ix.clear();
-        self.grammar_mut().override_config = Some(OverrideConfig {
+
+        let grammar = self
+            .grammar_mut()
+            .ok_or_else(|| anyhow!("cannot mutate grammar"))?;
+        grammar.override_config = Some(OverrideConfig {
             query,
             values: override_configs_by_id,
         });
@@ -1167,7 +1189,10 @@ impl Language {
     }
 
     pub fn with_redaction_query(mut self, source: &str) -> anyhow::Result<Self> {
-        let grammar = self.grammar_mut();
+        let grammar = self
+            .grammar_mut()
+            .ok_or_else(|| anyhow!("cannot mutate grammar"))?;
+
         let query = Query::new(&grammar.ts_language, source)?;
         let mut redaction_capture_ix = None;
         get_capture_indices(&query, &mut [("redact", &mut redaction_capture_ix)]);
@@ -1182,8 +1207,8 @@ impl Language {
         Ok(self)
     }
 
-    fn grammar_mut(&mut self) -> &mut Grammar {
-        Arc::get_mut(self.grammar.as_mut().unwrap()).unwrap()
+    fn grammar_mut(&mut self) -> Option<&mut Grammar> {
+        Arc::get_mut(self.grammar.as_mut()?)
     }
 
     pub fn name(&self) -> Arc<str> {

crates/language/src/language_registry.rs

@@ -431,15 +431,16 @@ impl LanguageRegistry {
                                 let language = async {
                                     let (config, queries) = (language.load)()?;
 
-                                    let grammar = if let Some(grammar) = config.grammar.clone() {
-                                        Some(this.get_or_load_grammar(grammar).await?)
+                                    if let Some(grammar) = config.grammar.clone() {
+                                        let grammar =
+                                            Some(this.get_or_load_grammar(grammar).await?);
+                                        Language::new_with_id(id, config, grammar)
+                                            .with_context_provider(provider)
+                                            .with_queries(queries)
                                     } else {
-                                        None
-                                    };
-
-                                    Language::new_with_id(id, config, grammar)
-                                        .with_context_provider(provider)
-                                        .with_queries(queries)
+                                        Ok(Language::new_with_id(id, config, None)
+                                            .with_context_provider(provider))
+                                    }
                                 }
                                 .await;