Avoid allocations in Picker file preview callback

The `FileLocation` and `PathOrId` types can borrow paths rather than
requiring them to be owned. This takes a refactor of the preview
functions and preview internals within `Picker`. With this change we
avoid an unnecessary `PathBuf` clone per render for any picker with a
file preview function (i.e. most pickers).

This refactor is not fully complete. The `PathOrId` is _sometimes_ an
owned `PathBuf`. This is for pragmatic reasons rather than technical
ones. We need a further refactor to introduce more core types like
`Location` in order to eliminate the Cow and only use `&Path`s within
`PathOrId`. This is left for future work as it will be a larger refactor
almost entirely fitting into the LSP commands module and helix-core -
i.e. mostly unrelated to refactoring the `Picker` code itself.

Co-authored-by: Pascal Kuthe <pascalkuthe@pm.me>

helix-term/src/commands.rs

@@ -2435,7 +2435,7 @@ fn global_search(cx: &mut Context) {
         },
     )
     .with_preview(|_editor, FileResult { path, line_num, .. }| {
-        Some((path.clone().into(), Some((*line_num, *line_num))))
+        Some((path.as_path().into(), Some((*line_num, *line_num))))
     })
     .with_history_register(Some(reg))
     .with_dynamic_query(get_files, Some(275));
@@ -3098,7 +3098,7 @@ fn changed_file_picker(cx: &mut Context) {
             }
         },
     )
-    .with_preview(|_editor, meta| Some((meta.path().to_path_buf().into(), None)));
+    .with_preview(|_editor, meta| Some((meta.path().into(), None)));
     let injector = picker.injector();
 
     cx.editor

helix-term/src/commands/dap.rs

@@ -61,7 +61,7 @@ fn thread_picker(
             .with_preview(move |editor, thread| {
                 let frames = editor.debugger.as_ref()?.stack_frames.get(&thread.id)?;
                 let frame = frames.first()?;
-                let path = frame.source.as_ref()?.path.clone()?;
+                let path = frame.source.as_ref()?.path.as_ref()?.as_path();
                 let pos = Some((
                     frame.line.saturating_sub(1),
                     frame.end_line.unwrap_or(frame.line).saturating_sub(1),
@@ -747,10 +747,10 @@ pub fn dap_switch_stack_frame(cx: &mut Context) {
         frame
             .source
             .as_ref()
-            .and_then(|source| source.path.clone())
+            .and_then(|source| source.path.as_ref())
             .map(|path| {
                 (
-                    path.into(),
+                    path.as_path().into(),
                     Some((
                         frame.line.saturating_sub(1),
                         frame.end_line.unwrap_or(frame.line).saturating_sub(1),

helix-term/src/commands/lsp.rs

@@ -64,6 +64,7 @@ macro_rules! language_server_with_feature {
 struct SymbolInformationItem {
     symbol: lsp::SymbolInformation,
     offset_encoding: OffsetEncoding,
+    uri: Uri,
 }
 
 struct DiagnosticStyles {
@@ -79,13 +80,10 @@ struct PickerDiagnostic {
     offset_encoding: OffsetEncoding,
 }
 
-fn location_to_file_location(location: &lsp::Location) -> FileLocation {
-    let path = location.uri.to_file_path().unwrap();
-    let line = Some((
-        location.range.start.line as usize,
-        location.range.end.line as usize,
-    ));
-    (path.into(), line)
+fn uri_to_file_location<'a>(uri: &'a Uri, range: &lsp::Range) -> Option<FileLocation<'a>> {
+    let path = uri.as_path()?;
+    let line = Some((range.start.line as usize, range.end.line as usize));
+    Some((path.into(), line))
 }
 
 fn jump_to_location(
@@ -278,7 +276,7 @@ fn diag_picker(
     )
     .with_preview(move |_editor, PickerDiagnostic { uri, diag, .. }| {
         let line = Some((diag.range.start.line as usize, diag.range.end.line as usize));
-        Some((uri.clone().as_path_buf()?.into(), line))
+        Some((uri.as_path()?.into(), line))
     })
     .truncate_start(false)
 }
@@ -287,6 +285,7 @@ pub fn symbol_picker(cx: &mut Context) {
     fn nested_to_flat(
         list: &mut Vec<SymbolInformationItem>,
         file: &lsp::TextDocumentIdentifier,
+        uri: &Uri,
         symbol: lsp::DocumentSymbol,
         offset_encoding: OffsetEncoding,
     ) {
@@ -301,9 +300,10 @@ pub fn symbol_picker(cx: &mut Context) {
                 container_name: None,
             },
             offset_encoding,
+            uri: uri.clone(),
         });
         for child in symbol.children.into_iter().flatten() {
-            nested_to_flat(list, file, child, offset_encoding);
+            nested_to_flat(list, file, uri, child, offset_encoding);
         }
     }
     let doc = doc!(cx.editor);
@@ -317,6 +317,9 @@ pub fn symbol_picker(cx: &mut Context) {
             let request = language_server.document_symbols(doc.identifier()).unwrap();
             let offset_encoding = language_server.offset_encoding();
             let doc_id = doc.identifier();
+            let doc_uri = doc
+                .uri()
+                .expect("docs with active language servers must be backed by paths");
 
             async move {
                 let json = request.await?;
@@ -331,6 +334,7 @@ pub fn symbol_picker(cx: &mut Context) {
                     lsp::DocumentSymbolResponse::Flat(symbols) => symbols
                         .into_iter()
                         .map(|symbol| SymbolInformationItem {
+                            uri: doc_uri.clone(),
                             symbol,
                             offset_encoding,
                         })
@@ -338,7 +342,13 @@ pub fn symbol_picker(cx: &mut Context) {
                     lsp::DocumentSymbolResponse::Nested(symbols) => {
                         let mut flat_symbols = Vec::new();
                         for symbol in symbols {
-                            nested_to_flat(&mut flat_symbols, &doc_id, symbol, offset_encoding)
+                            nested_to_flat(
+                                &mut flat_symbols,
+                                &doc_id,
+                                &doc_uri,
+                                symbol,
+                                offset_encoding,
+                            )
                         }
                         flat_symbols
                     }
@@ -388,7 +398,7 @@ pub fn symbol_picker(cx: &mut Context) {
                 },
             )
             .with_preview(move |_editor, item| {
-                Some(location_to_file_location(&item.symbol.location))
+                uri_to_file_location(&item.uri, &item.symbol.location.range)
             })
             .truncate_start(false);
 
@@ -431,9 +441,19 @@ pub fn workspace_symbol_picker(cx: &mut Context) {
                         serde_json::from_value::<Option<Vec<lsp::SymbolInformation>>>(json)?
                             .unwrap_or_default()
                             .into_iter()
-                            .map(|symbol| SymbolInformationItem {
-                                symbol,
-                                offset_encoding,
+                            .filter_map(|symbol| {
+                                let uri = match Uri::try_from(&symbol.location.uri) {
+                                    Ok(uri) => uri,
+                                    Err(err) => {
+                                        log::warn!("discarding symbol with invalid URI: {err}");
+                                        return None;
+                                    }
+                                };
+                                Some(SymbolInformationItem {
+                                    symbol,
+                                    uri,
+                                    offset_encoding,
+                                })
                             })
                             .collect();
 
@@ -467,15 +487,11 @@ pub fn workspace_symbol_picker(cx: &mut Context) {
         })
         .without_filtering(),
         ui::PickerColumn::new("path", |item: &SymbolInformationItem, _| {
-            if let Ok(uri) = Uri::try_from(&item.symbol.location.uri) {
-                if let Some(path) = uri.as_path() {
-                    path::get_relative_path(path)
-                        .to_string_lossy()
-                        .to_string()
-                        .into()
-                } else {
-                    item.symbol.location.uri.to_string().into()
-                }
+            if let Some(path) = item.uri.as_path() {
+                path::get_relative_path(path)
+                    .to_string_lossy()
+                    .to_string()
+                    .into()
             } else {
                 item.symbol.location.uri.to_string().into()
             }
@@ -496,7 +512,7 @@ pub fn workspace_symbol_picker(cx: &mut Context) {
             );
         },
     )
-    .with_preview(|_editor, item| Some(location_to_file_location(&item.symbol.location)))
+    .with_preview(|_editor, item| uri_to_file_location(&item.uri, &item.symbol.location.range))
     .with_dynamic_query(get_symbols, None)
     .truncate_start(false);
 
@@ -875,7 +891,31 @@ fn goto_impl(
             let picker = Picker::new(columns, 0, locations, cwdir, move |cx, location, action| {
                 jump_to_location(cx.editor, location, offset_encoding, action)
             })
-            .with_preview(move |_editor, location| Some(location_to_file_location(location)));
+            .with_preview(move |_editor, location| {
+                use crate::ui::picker::PathOrId;
+
+                let lines = Some((
+                    location.range.start.line as usize,
+                    location.range.end.line as usize,
+                ));
+
+                // TODO: we should avoid allocating by doing the Uri conversion ahead of time.
+                //
+                // To do this, introduce a `Location` type in `helix-core` that reuses the core
+                // `Uri` type instead of the LSP `Url` type and replaces the LSP `Range` type.
+                // Refactor the callers of `goto_impl` to pass iterators that translate the
+                // LSP location type to the custom one in core, or have them collect and pass
+                // `Vec<Location>`s. Replace the `uri_to_file_location` function with
+                // `location_to_file_location` that takes only `&helix_core::Location` as
+                // parameters.
+                //
+                // By doing this we can also eliminate the duplicated URI info in the
+                // `SymbolInformationItem` type and introduce a custom Symbol type in `helix-core`
+                // which will be reused in the future for tree-sitter based symbol pickers.
+                let path = Uri::try_from(&location.uri).ok()?.as_path_buf()?;
+                #[allow(deprecated)]
+                Some((PathOrId::from_path_buf(path), lines))
+            });
             compositor.push(Box::new(overlaid(picker)));
         }
     }

helix-term/src/ui/mod.rs

@@ -244,7 +244,7 @@ pub fn file_picker(root: PathBuf, config: &helix_view::editor::Config) -> FilePi
             }
         },
     )
-    .with_preview(|_editor, path| Some((path.clone().into(), None)));
+    .with_preview(|_editor, path| Some((path.as_path().into(), None)));
     let injector = picker.injector();
     let timeout = std::time::Instant::now() + std::time::Duration::from_millis(30);
 

helix-term/src/ui/picker.rs

@@ -60,37 +60,41 @@ pub const MIN_AREA_WIDTH_FOR_PREVIEW: u16 = 72;
 pub const MAX_FILE_SIZE_FOR_PREVIEW: u64 = 10 * 1024 * 1024;
 
 #[derive(PartialEq, Eq, Hash)]
-pub enum PathOrId {
+pub enum PathOrId<'a> {
     Id(DocumentId),
-    Path(Arc<Path>),
+    // See [PathOrId::from_path_buf]: this will eventually become `Path(&Path)`.
+    Path(Cow<'a, Path>),
 }
 
-impl PathOrId {
-    fn get_canonicalized(self) -> Self {
-        use PathOrId::*;
-        match self {
-            Path(path) => Path(helix_stdx::path::canonicalize(&path).into()),
-            Id(id) => Id(id),
-        }
+impl<'a> PathOrId<'a> {
+    /// Creates a [PathOrId] from a PathBuf
+    ///
+    /// # Deprecated
+    /// The owned version of PathOrId will be removed in a future refactor
+    /// and replaced with `&'a Path`. See the caller of this function for
+    /// more details on its removal.
+    #[deprecated]
+    pub fn from_path_buf(path_buf: PathBuf) -> Self {
+        Self::Path(Cow::Owned(path_buf))
     }
 }
 
-impl From<PathBuf> for PathOrId {
-    fn from(v: PathBuf) -> Self {
-        Self::Path(v.as_path().into())
+impl<'a> From<&'a Path> for PathOrId<'a> {
+    fn from(path: &'a Path) -> Self {
+        Self::Path(Cow::Borrowed(path))
     }
 }
 
-impl From<DocumentId> for PathOrId {
+impl<'a> From<DocumentId> for PathOrId<'a> {
     fn from(v: DocumentId) -> Self {
         Self::Id(v)
     }
 }
 
-type FileCallback<T> = Box<dyn Fn(&Editor, &T) -> Option<FileLocation>>;
+type FileCallback<T> = Box<dyn for<'a> Fn(&'a Editor, &'a T) -> Option<FileLocation<'a>>>;
 
 /// File path and range of lines (used to align and highlight lines)
-pub type FileLocation = (PathOrId, Option<(usize, usize)>);
+pub type FileLocation<'a> = (PathOrId<'a>, Option<(usize, usize)>);
 
 pub enum CachedPreview {
     Document(Box<Document>),
@@ -400,7 +404,7 @@ impl<T: 'static + Send + Sync, D: 'static + Send + Sync> Picker<T, D> {
 
     pub fn with_preview(
         mut self,
-        preview_fn: impl Fn(&Editor, &T) -> Option<FileLocation> + 'static,
+        preview_fn: impl for<'a> Fn(&'a Editor, &'a T) -> Option<FileLocation<'a>> + 'static,
     ) -> Self {
         self.file_fn = Some(Box::new(preview_fn));
         // assumption: if we have a preview we are matching paths... If this is ever
@@ -544,40 +548,35 @@ impl<T: 'static + Send + Sync, D: 'static + Send + Sync> Picker<T, D> {
         }
     }
 
-    fn current_file(&self, editor: &Editor) -> Option<FileLocation> {
-        self.selection()
-            .and_then(|current| (self.file_fn.as_ref()?)(editor, current))
-            .map(|(path_or_id, line)| (path_or_id.get_canonicalized(), line))
-    }
-
-    /// Get (cached) preview for a given path. If a document corresponding
+    /// Get (cached) preview for the currently selected item. If a document corresponding
     /// to the path is already open in the editor, it is used instead.
     fn get_preview<'picker, 'editor>(
         &'picker mut self,
-        path_or_id: PathOrId,
         editor: &'editor Editor,
-    ) -> Preview<'picker, 'editor> {
+    ) -> Option<(Preview<'picker, 'editor>, Option<(usize, usize)>)> {
+        let current = self.selection()?;
+        let (path_or_id, range) = (self.file_fn.as_ref()?)(editor, current)?;
+
         match path_or_id {
             PathOrId::Path(path) => {
-                if let Some(doc) = editor.document_by_path(&path) {
-                    return Preview::EditorDocument(doc);
+                let path = path.as_ref();
+                if let Some(doc) = editor.document_by_path(path) {
+                    return Some((Preview::EditorDocument(doc), range));
                 }
 
-                if self.preview_cache.contains_key(&path) {
-                    let preview = &self.preview_cache[&path];
-                    match preview {
-                        // If the document isn't highlighted yet, attempt to highlight it.
-                        CachedPreview::Document(doc) if doc.language_config().is_none() => {
-                            helix_event::send_blocking(
-                                &self.preview_highlight_handler,
-                                path.clone(),
-                            );
-                        }
-                        _ => (),
+                if self.preview_cache.contains_key(path) {
+                    // NOTE: we use `HashMap::get_key_value` here instead of indexing so we can
+                    // retrieve the `Arc<Path>` key. The `path` in scope here is a `&Path` and
+                    // we can cheaply clone the key for the preview highlight handler.
+                    let (path, preview) = self.preview_cache.get_key_value(path).unwrap();
+                    if matches!(preview, CachedPreview::Document(doc) if doc.language_config().is_none())
+                    {
+                        helix_event::send_blocking(&self.preview_highlight_handler, path.clone());
                     }
-                    return Preview::Cached(preview);
+                    return Some((Preview::Cached(preview), range));
                 }
 
+                let path: Arc<Path> = path.into();
                 let data = std::fs::File::open(&path).and_then(|file| {
                     let metadata = file.metadata()?;
                     // Read up to 1kb to detect the content type
@@ -607,11 +606,11 @@ impl<T: 'static + Send + Sync, D: 'static + Send + Sync> Picker<T, D> {
                     )
                     .unwrap_or(CachedPreview::NotFound);
                 self.preview_cache.insert(path.clone(), preview);
-                Preview::Cached(&self.preview_cache[&path])
+                Some((Preview::Cached(&self.preview_cache[&path]), range))
             }
             PathOrId::Id(id) => {
                 let doc = editor.documents.get(&id).unwrap();
-                Preview::EditorDocument(doc)
+                Some((Preview::EditorDocument(doc), range))
             }
         }
     }
@@ -816,8 +815,7 @@ impl<T: 'static + Send + Sync, D: 'static + Send + Sync> Picker<T, D> {
         let inner = inner.inner(margin);
         BLOCK.render(area, surface);
 
-        if let Some((path, range)) = self.current_file(cx.editor) {
-            let preview = self.get_preview(path, cx.editor);
+        if let Some((preview, range)) = self.get_preview(cx.editor) {
             let doc = match preview.document() {
                 Some(doc)
                     if range.map_or(true, |(start, end)| {