[perf] use StackVector to save memory & CPU

Author: atuchin-m

src/filters/network.rs

@@ -15,9 +15,7 @@ use crate::filters::abstract_network::{
 use crate::lists::ParseOptions;
 use crate::regex_manager::RegexManager;
 use crate::request;
-use crate::utils::{self, Hash};
-
-pub(crate) const TOKENS_BUFFER_SIZE: usize = 200;
+use crate::utils::{self, Hash, TokensBuffer};
 
 /// For now, only support `$removeparam` with simple alphanumeric/dash/underscore patterns.
 static VALID_PARAM: Lazy<Regex> = Lazy::new(|| Regex::new(r"^[a-zA-Z0-9_\-]+$").unwrap());
@@ -895,7 +893,7 @@ impl NetworkFilter {
     }
 
     pub fn get_tokens_optimized(&self) -> FilterTokens {
-        let mut tokens: Vec<Hash> = Vec::with_capacity(TOKENS_BUFFER_SIZE);
+        let mut tokens = TokensBuffer::default();
 
         // If there is only one domain and no domain negation, we also use this
         // domain as a token.
@@ -905,7 +903,7 @@ impl NetworkFilter {
         {
             if let Some(domains) = self.opt_domains.as_ref() {
                 if let Some(domain) = domains.first() {
-                    tokens.push(*domain)
+                    tokens.push(*domain);
                 }
             }
         }
@@ -963,10 +961,7 @@ impl NetworkFilter {
                 tokens.push(utils::fast_hash("https"));
             }
 
-            // Remake a vector to drop extra capacity.
-            let mut t = Vec::with_capacity(tokens.len());
-            t.extend(tokens);
-            FilterTokens::Other(t)
+            FilterTokens::Other(tokens.into_vec())
         }
     }
 }

src/flatbuffers/unsafe_tools.rs

@@ -1,6 +1,7 @@
 //! Unsafe utility functions for working with flatbuffers and other low-level operations.
 
 use crate::filters::flatbuffer_generated::fb;
+use std::mem::MaybeUninit;
 
 // Minimum alignment for the beginning of the flatbuffer data.
 const MIN_ALIGNMENT: usize = 8;
@@ -97,3 +98,46 @@ impl VerifiedFlatbufferMemory {
         &self.raw_data[self.start..]
     }
 }
+
+/// A simple stack-allocated vector.
+/// It is used to avoid allocations when the vector is small.
+pub(crate) struct StackVector<T, const MAX_SIZE: usize> {
+    data: [MaybeUninit<T>; MAX_SIZE],
+    size: usize,
+}
+
+impl<T, const MAX_SIZE: usize> Default for StackVector<T, MAX_SIZE>
+where
+    T: Default + Copy,
+{
+    fn default() -> Self {
+        Self {
+            data: [MaybeUninit::uninit(); MAX_SIZE],
+            size: 0,
+        }
+    }
+}
+
+impl<T, const MAX_SIZE: usize> StackVector<T, MAX_SIZE> {
+    pub fn push(&mut self, value: T) -> bool {
+        if self.size < MAX_SIZE {
+            self.data[self.size] = MaybeUninit::new(value);
+            self.size += 1;
+            true
+        } else {
+            false
+        }
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.size == 0
+    }
+
+    pub fn into_vec(self) -> Vec<T> {
+        let mut v = Vec::with_capacity(self.size);
+        for i in 0..self.size {
+            v.push(unsafe { self.data[i].assume_init_read() });
+        }
+        v
+    }
+}

src/request.rs

@@ -239,11 +239,11 @@ impl Request {
 }
 
 fn calculate_tokens(url_lower_cased: &str) -> Vec<utils::Hash> {
-    let mut tokens = vec![];
+    let mut tokens = utils::TokensBuffer::default();
     utils::tokenize_pooled(url_lower_cased, &mut tokens);
     // Add zero token as a fallback to wildcard rule bucket
     tokens.push(0);
-    tokens
+    tokens.into_vec()
 }
 
 #[cfg(test)]

src/utils.rs

@@ -6,6 +6,8 @@ use seahash::hash;
 #[cfg(target_pointer_width = "32")]
 use seahash::reference::hash;
 
+use crate::flatbuffers::unsafe_tools::StackVector;
+
 pub type Hash = u64;
 
 // A smaller version of Hash that is used in serialized format.
@@ -27,26 +29,21 @@ fn is_allowed_filter(ch: char) -> bool {
     ch.is_alphanumeric() || ch == '%'
 }
 
-pub(crate) const TOKENS_BUFFER_SIZE: usize = 128;
-pub(crate) const TOKENS_BUFFER_RESERVED: usize = 1;
-const TOKENS_MAX: usize = TOKENS_BUFFER_SIZE - TOKENS_BUFFER_RESERVED;
+pub(crate) type TokensBuffer = StackVector<Hash, 200>;
 
 fn fast_tokenizer_no_regex(
     pattern: &str,
     is_allowed_code: &dyn Fn(char) -> bool,
     skip_first_token: bool,
     skip_last_token: bool,
-    tokens_buffer: &mut Vec<Hash>,
+    tokens_buffer: &mut TokensBuffer,
 ) {
     // let mut tokens_buffer_index = 0;
     let mut inside: bool = false;
     let mut start = 0;
     let mut preceding_ch: Option<char> = None; // Used to check if a '*' is not just before a token
 
     for (i, c) in pattern.char_indices() {
-        if tokens_buffer.len() >= TOKENS_MAX {
-            return;
-        }
         if is_allowed_code(c) {
             if !inside {
                 inside = true;
@@ -75,17 +72,17 @@ fn fast_tokenizer_no_regex(
     }
 }
 
-pub(crate) fn tokenize_pooled(pattern: &str, tokens_buffer: &mut Vec<Hash>) {
+pub(crate) fn tokenize_pooled(pattern: &str, tokens_buffer: &mut TokensBuffer) {
     fast_tokenizer_no_regex(pattern, &is_allowed_filter, false, false, tokens_buffer);
 }
 
 pub fn tokenize(pattern: &str) -> Vec<Hash> {
-    let mut tokens_buffer: Vec<Hash> = Vec::with_capacity(TOKENS_BUFFER_SIZE);
+    let mut tokens_buffer = TokensBuffer::default();
     tokenize_to(pattern, &mut tokens_buffer);
-    tokens_buffer
+    tokens_buffer.into_vec()
 }
 
-pub(crate) fn tokenize_to(pattern: &str, tokens_buffer: &mut Vec<Hash>) {
+pub(crate) fn tokenize_to(pattern: &str, tokens_buffer: &mut TokensBuffer) {
     fast_tokenizer_no_regex(pattern, &is_allowed_filter, false, false, tokens_buffer);
 }
 
@@ -95,21 +92,21 @@ pub(crate) fn tokenize_filter(
     skip_first_token: bool,
     skip_last_token: bool,
 ) -> Vec<Hash> {
-    let mut tokens_buffer: Vec<Hash> = Vec::with_capacity(TOKENS_BUFFER_SIZE);
+    let mut tokens_buffer = TokensBuffer::default();
     tokenize_filter_to(
         pattern,
         skip_first_token,
         skip_last_token,
         &mut tokens_buffer,
     );
-    tokens_buffer
+    tokens_buffer.into_vec()
 }
 
 pub(crate) fn tokenize_filter_to(
     pattern: &str,
     skip_first_token: bool,
     skip_last_token: bool,
-    tokens_buffer: &mut Vec<Hash>,
+    tokens_buffer: &mut TokensBuffer,
 ) {
     fast_tokenizer_no_regex(
         pattern,