Split get_document_or_error() from get_document() (#313)

Author: DavisVaughan

crates/lsp/src/handlers_ext.rs

@@ -23,7 +23,7 @@ pub(crate) enum ViewFileKind {
 }
 
 pub(crate) fn view_file(params: ViewFileParams, state: &WorldState) -> anyhow::Result<String> {
-    let doc = state.get_document(&params.text_document.uri)?;
+    let doc = state.get_document_or_error(&params.text_document.uri)?;
 
     match params.kind {
         ViewFileKind::TreeSitter => {

crates/lsp/src/handlers_format.rs

@@ -23,7 +23,7 @@ pub(crate) fn document_formatting(
     state: &WorldState,
 ) -> anyhow::Result<Option<Vec<lsp_types::TextEdit>>> {
     let uri = &params.text_document.uri;
-    let doc = state.get_document(uri)?;
+    let doc = state.get_document_or_error(uri)?;
 
     let workspace_settings = lsp_state.workspace_document_settings(uri);
 
@@ -68,7 +68,7 @@ pub(crate) fn document_range_formatting(
     state: &WorldState,
 ) -> anyhow::Result<Option<Vec<lsp_types::TextEdit>>> {
     let uri = &params.text_document.uri;
-    let doc = state.get_document(uri)?;
+    let doc = state.get_document_or_error(uri)?;
 
     let workspace_settings = lsp_state.workspace_document_settings(uri);
 

crates/lsp/src/handlers_state.rs

@@ -167,7 +167,7 @@ pub(crate) fn did_change(
     state: &mut WorldState,
 ) -> anyhow::Result<()> {
     let uri = &params.text_document.uri;
-    let doc = state.get_document_mut(uri)?;
+    let doc = state.get_document_mut_or_error(uri)?;
     doc.on_did_change(params);
 
     Ok(())
@@ -249,7 +249,7 @@ pub(crate) fn did_change_formatting_options(
     opts: &FormattingOptions,
     state: &mut WorldState,
 ) {
-    let Ok(doc) = state.get_document_mut(uri) else {
+    let Some(doc) = state.get_document_mut(uri) else {
         return;
     };
 
@@ -400,7 +400,7 @@ fn update_documents_config(
         let config: DocumentSettings = config.into();
 
         // Finally, update the document's config
-        state.get_document_mut(&uri)?.settings = config;
+        state.get_document_mut_or_error(&uri)?.settings = config;
     }
 
     Ok(())

crates/lsp/src/state.rs

@@ -41,19 +41,25 @@ pub(crate) struct WorldState {
 }
 
 impl WorldState {
-    pub(crate) fn get_document(&self, uri: &Url) -> anyhow::Result<&Document> {
-        if let Some(doc) = self.documents.get(uri) {
-            Ok(doc)
-        } else {
-            Err(anyhow!("Can't find document for URI {uri}"))
+    pub(crate) fn get_document(&self, uri: &Url) -> Option<&Document> {
+        self.documents.get(uri)
+    }
+
+    pub(crate) fn get_document_mut(&mut self, uri: &Url) -> Option<&mut Document> {
+        self.documents.get_mut(uri)
+    }
+
+    pub(crate) fn get_document_or_error(&self, uri: &Url) -> anyhow::Result<&Document> {
+        match self.get_document(uri) {
+            Some(doc) => Ok(doc),
+            None => Err(anyhow!("Can't find document for URI {uri}")),
         }
     }
 
-    pub(crate) fn get_document_mut(&mut self, uri: &Url) -> anyhow::Result<&mut Document> {
-        if let Some(doc) = self.documents.get_mut(uri) {
-            Ok(doc)
-        } else {
-            Err(anyhow!("Can't find document for URI {uri}"))
+    pub(crate) fn get_document_mut_or_error(&mut self, uri: &Url) -> anyhow::Result<&mut Document> {
+        match self.get_document_mut(uri) {
+            Some(doc) => Ok(doc),
+            None => Err(anyhow!("Can't find document for URI {uri}")),
         }
     }