From 32f69dbe156030de5ae589d968efc2825bd0485f Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Wed, 5 Aug 2020 09:06:00 -0700
Subject: [PATCH] tags, highlight: Limit the size of buffers that are retained
 in memory

---
 highlight/src/lib.rs | 16 ++++++++++++----
 tags/src/c_lib.rs    | 12 ++++++++++--
 2 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/highlight/src/lib.rs b/highlight/src/lib.rs
index bb110219..1cffefa2 100644
--- a/highlight/src/lib.rs
+++ b/highlight/src/lib.rs
@@ -10,6 +10,8 @@ use tree_sitter::{
 };
 
 const CANCELLATION_CHECK_INTERVAL: usize = 100;
+const BUFFER_HTML_RESERVE_CAPACITY: usize = 10 * 1024;
+const BUFFER_LINES_RESERVE_CAPACITY: usize = 1000;
 
 /// Indicates which highlight should be applied to a region of source code.
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
@@ -884,11 +886,13 @@ where
 
 impl HtmlRenderer {
     pub fn new() -> Self {
-        HtmlRenderer {
-            html: Vec::new(),
-            line_offsets: vec![0],
+        let mut result = HtmlRenderer {
+            html: Vec::with_capacity(BUFFER_HTML_RESERVE_CAPACITY),
+            line_offsets: Vec::with_capacity(BUFFER_LINES_RESERVE_CAPACITY),
             carriage_return_highlight: None,
-        }
+        };
+        result.line_offsets.push(0);
+        result
     }
 
     pub fn set_carriage_return_highlight(&mut self, highlight: Option<Highlight>) {
@@ -896,6 +900,10 @@ impl HtmlRenderer {
     }
 
     pub fn reset(&mut self) {
+        self.html.truncate(BUFFER_HTML_RESERVE_CAPACITY);
+        self.line_offsets.truncate(BUFFER_LINES_RESERVE_CAPACITY);
+        self.html.shrink_to_fit();
+        self.line_offsets.shrink_to_fit();
         self.html.clear();
         self.line_offsets.clear();
         self.line_offsets.push(0);
diff --git a/tags/src/c_lib.rs b/tags/src/c_lib.rs
index 07e1e19a..c2bec6ca 100644
--- a/tags/src/c_lib.rs
+++ b/tags/src/c_lib.rs
@@ -6,6 +6,9 @@ use std::sync::atomic::AtomicUsize;
 use std::{fmt, slice, str};
 use tree_sitter::Language;
 
+const BUFFER_TAGS_RESERVE_CAPACITY: usize = 100;
+const BUFFER_DOCS_RESERVE_CAPACITY: usize = 1024;
+
 #[repr(C)]
 #[derive(Debug, PartialEq, Eq)]
 pub enum TSTagsError {
@@ -116,8 +119,13 @@ pub extern "C" fn ts_tagger_tag(
     let scope_name = unsafe { unwrap(CStr::from_ptr(scope_name).to_str()) };
 
     if let Some(config) = tagger.languages.get(scope_name) {
+        buffer.tags.truncate(BUFFER_TAGS_RESERVE_CAPACITY);
+        buffer.docs.truncate(BUFFER_DOCS_RESERVE_CAPACITY);
+        buffer.tags.shrink_to_fit();
+        buffer.docs.shrink_to_fit();
         buffer.tags.clear();
         buffer.docs.clear();
+
         let source_code = unsafe { slice::from_raw_parts(source_code, source_code_len as usize) };
         let cancellation_flag = unsafe { cancellation_flag.as_ref() };
 
@@ -182,8 +190,8 @@ pub extern "C" fn ts_tagger_tag(
 pub extern "C" fn ts_tags_buffer_new() -> *mut TSTagsBuffer {
     Box::into_raw(Box::new(TSTagsBuffer {
         context: TagsContext::new(),
-        tags: Vec::with_capacity(64),
-        docs: Vec::with_capacity(64),
+        tags: Vec::with_capacity(BUFFER_TAGS_RESERVE_CAPACITY),
+        docs: Vec::with_capacity(BUFFER_DOCS_RESERVE_CAPACITY),
     }))
 }