From 6992c19aee80ad79e858cbb504a0ed1a04eab2dd Mon Sep 17 00:00:00 2001
From: Antonio Scandurra <me@as-cii.com>
Date: Fri, 28 May 2021 11:48:49 +0200
Subject: [PATCH] Acquire and release `QueryCursor`s via a `QueryCursorHandle`

---
 zed/src/editor/buffer/mod.rs | 64 ++++++++++++++++++++++--------------
 1 file changed, 40 insertions(+), 24 deletions(-)

diff --git a/zed/src/editor/buffer/mod.rs b/zed/src/editor/buffer/mod.rs
index aa7a9f491372e2a2194bacd57ed59009043c0d7c..81272ced2d7e2eb0c8c5099c1241acd406e08cd5 100644
--- a/zed/src/editor/buffer/mod.rs
+++ b/zed/src/editor/buffer/mod.rs
@@ -30,7 +30,7 @@ use std::{
     hash::BuildHasher,
     iter::{self, Iterator},
     mem,
-    ops::Range,
+    ops::{Deref, DerefMut, Range},
     str,
     sync::Arc,
     time::{Duration, Instant, SystemTime, UNIX_EPOCH},
@@ -69,6 +69,42 @@ lazy_static! {
     static ref QUERY_CURSORS: Mutex<Vec<QueryCursor>> = Default::default();
 }
 
+struct QueryCursorHandle(Option<QueryCursor>);
+
+impl QueryCursorHandle {
+    fn new() -> Self {
+        QueryCursorHandle(Some(
+            QUERY_CURSORS
+                .lock()
+                .pop()
+                .unwrap_or_else(|| QueryCursor::new()),
+        ))
+    }
+}
+
+impl Deref for QueryCursorHandle {
+    type Target = QueryCursor;
+
+    fn deref(&self) -> &Self::Target {
+        self.0.as_ref().unwrap()
+    }
+}
+
+impl DerefMut for QueryCursorHandle {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        self.0.as_mut().unwrap()
+    }
+}
+
+impl Drop for QueryCursorHandle {
+    fn drop(&mut self) {
+        let mut cursor = self.0.take().unwrap();
+        cursor.set_byte_range(0, usize::MAX);
+        cursor.set_point_range(Point::zero().into(), Point::MAX.into());
+        QUERY_CURSORS.lock().push(cursor)
+    }
+}
+
 pub struct Buffer {
     fragments: SumTree<Fragment>,
     visible_text: Rope,
@@ -520,7 +556,7 @@ impl Buffer {
             text: self.visible_text.clone(),
             tree: self.syntax_tree(),
             language: self.language.clone(),
-            query_cursor: Some(acquire_query_cursor()),
+            query_cursor: QueryCursorHandle::new(),
         }
     }
 
@@ -2046,7 +2082,7 @@ pub struct Snapshot {
     text: Rope,
     tree: Option<Tree>,
     language: Option<Arc<Language>>,
-    query_cursor: Option<QueryCursor>,
+    query_cursor: QueryCursorHandle,
 }
 
 impl Snapshot {
@@ -2065,8 +2101,7 @@ impl Snapshot {
     pub fn highlighted_text_for_range(&mut self, range: Range<usize>) -> HighlightedChunks {
         let chunks = self.text.chunks_in_range(range.clone());
         if let Some((language, tree)) = self.language.as_ref().zip(self.tree.as_ref()) {
-            let query_cursor = self.query_cursor.as_mut().unwrap();
-            let mut captures = query_cursor.captures(
+            let mut captures = self.query_cursor.captures(
                 &language.highlight_query,
                 tree.root_node(),
                 TextProvider(&self.text),
@@ -2109,25 +2144,6 @@ impl Snapshot {
     }
 }
 
-impl Drop for Snapshot {
-    fn drop(&mut self) {
-        release_query_cursor(self.query_cursor.take().unwrap());
-    }
-}
-
-fn acquire_query_cursor() -> QueryCursor {
-    QUERY_CURSORS
-        .lock()
-        .pop()
-        .unwrap_or_else(|| QueryCursor::new())
-}
-
-fn release_query_cursor(mut cursor: QueryCursor) {
-    cursor.set_byte_range(0, usize::MAX);
-    cursor.set_point_range(Point::zero().into(), Point::MAX.into());
-    QUERY_CURSORS.lock().push(cursor)
-}
-
 struct RopeBuilder<'a> {
     old_visible_cursor: rope::Cursor<'a>,
     old_deleted_cursor: rope::Cursor<'a>,