Change get_tokens_optimized()

Author: atuchin-m

src/filters/fb_network_builder.rs

@@ -7,6 +7,7 @@ use flatbuffers::WIPOffset;
 use crate::filters::fb_builder::EngineFlatBuilder;
 use crate::filters::network::{FilterTokens, NetworkFilter};
 use crate::filters::token_selector::TokenSelector;
+use crate::utils::TokensBuffer;
 
 use crate::filters::network::NetworkFilterMaskHelper;
 use crate::flatbuffers::containers::flat_multimap::FlatMultiMapBuilder;
@@ -134,6 +135,7 @@ impl<'a> FlatSerialize<'a, EngineFlatBuilder<'a>> for NetworkFilterListBuilder {
         let mut optimizable = HashMap::<ShortHash, Vec<NetworkFilter>>::new();
 
         let mut token_frequencies = TokenSelector::new(rule_list.filters.len());
+        let mut tokens_buffer = TokensBuffer::default();
 
         {
             for network_filter in rule_list.filters {
@@ -157,7 +159,7 @@ impl<'a> FlatSerialize<'a, EngineFlatBuilder<'a>> for NetworkFilterListBuilder {
                     }
                 };
 
-                let multi_tokens = network_filter.get_tokens_optimized();
+                let multi_tokens = network_filter.get_tokens_optimized(&mut tokens_buffer);
                 match multi_tokens {
                     FilterTokens::Empty => {
                         // No tokens, add to fallback bucket (token 0)

src/filters/network.rs

@@ -310,10 +310,10 @@ pub enum FilterPart {
 }
 
 #[derive(Debug, PartialEq)]
-pub enum FilterTokens {
+pub enum FilterTokens<'a> {
     Empty,
-    OptDomains(Vec<Hash>),
-    Other(Vec<Hash>),
+    OptDomains(&'a [Hash]),
+    Other(&'a [Hash]),
 }
 
 pub struct FilterPartIterator<'a> {
@@ -883,17 +883,18 @@ impl NetworkFilter {
 
     #[deprecated(since = "0.11.1", note = "use get_tokens_optimized instead")]
     pub fn get_tokens(&self) -> Vec<Vec<Hash>> {
-        match self.get_tokens_optimized() {
+        let mut tokens_buffer = TokensBuffer::default();
+        match self.get_tokens_optimized(&mut tokens_buffer) {
             FilterTokens::OptDomains(domains) => {
-                domains.into_iter().map(|domain| vec![domain]).collect()
+                domains.into_iter().map(|domain| vec![*domain]).collect()
             }
-            FilterTokens::Other(tokens) => vec![tokens],
+            FilterTokens::Other(tokens) => vec![tokens.to_vec()],
             FilterTokens::Empty => vec![],
         }
     }
 
-    pub fn get_tokens_optimized(&self) -> FilterTokens {
-        let mut tokens = TokensBuffer::default();
+    pub fn get_tokens_optimized<'a>(&'a self, tokens_buffer: &'a mut TokensBuffer) -> FilterTokens<'a> {
+        tokens_buffer.clear();
 
         // If there is only one domain and no domain negation, we also use this
         // domain as a token.
@@ -903,7 +904,7 @@ impl NetworkFilter {
         {
             if let Some(domains) = self.opt_domains.as_ref() {
                 if let Some(domain) = domains.first() {
-                    tokens.push(*domain);
+                    tokens_buffer.push(*domain);
                 }
             }
         }
@@ -916,7 +917,7 @@ impl NetworkFilter {
                         (self.is_plain() || self.is_regex()) && !self.is_right_anchor();
                     let skip_first_token = self.is_right_anchor();
 
-                    utils::tokenize_filter_to(f, skip_first_token, skip_last_token, &mut tokens);
+                    utils::tokenize_filter_to(f, skip_first_token, skip_last_token, tokens_buffer);
                 }
             }
             FilterPart::AnyOf(_) => (), // across AnyOf set of filters no single token is guaranteed to match to a request
@@ -926,42 +927,42 @@ impl NetworkFilter {
         // Append tokens from hostname, if any
         if !self.mask.contains(NetworkFilterMask::IS_HOSTNAME_REGEX) {
             if let Some(hostname) = self.hostname.as_ref() {
-                utils::tokenize_to(hostname, &mut tokens);
+                utils::tokenize_to(hostname, tokens_buffer);
             }
         } else if let Some(hostname) = self.hostname.as_ref() {
             // Find last dot to tokenize the prefix
             let last_dot_pos = hostname.rfind('.');
             if let Some(last_dot_pos) = last_dot_pos {
-                utils::tokenize_to(&hostname[..last_dot_pos], &mut tokens);
+                utils::tokenize_to(&hostname[..last_dot_pos], tokens_buffer);
             }
         }
 
-        if tokens.is_empty() && self.mask.contains(NetworkFilterMask::IS_REMOVEPARAM) {
+        if tokens_buffer.is_empty() && self.mask.contains(NetworkFilterMask::IS_REMOVEPARAM) {
             if let Some(removeparam) = &self.modifier_option {
                 if VALID_PARAM.is_match(removeparam) {
-                    utils::tokenize_to(&removeparam.to_ascii_lowercase(), &mut tokens);
+                    utils::tokenize_to(&removeparam.to_ascii_lowercase(), tokens_buffer);
                 }
             }
         }
 
         // If we got no tokens for the filter/hostname part, then we will dispatch
         // this filter in multiple buckets based on the domains option.
-        if tokens.is_empty() && self.opt_domains.is_some() && self.opt_not_domains.is_none() {
+        if tokens_buffer.is_empty() && self.opt_domains.is_some() && self.opt_not_domains.is_none() {
             if let Some(opt_domains) = self.opt_domains.as_ref() {
                 if !opt_domains.is_empty() {
-                    return FilterTokens::OptDomains(opt_domains.clone());
+                    return FilterTokens::OptDomains(opt_domains);
                 }
             }
             FilterTokens::Empty
         } else {
             // Add optional token for protocol
             if self.for_http() && !self.for_https() {
-                tokens.push(utils::fast_hash("http"));
+                tokens_buffer.push(utils::fast_hash("http"));
             } else if self.for_https() && !self.for_http() {
-                tokens.push(utils::fast_hash("https"));
+                tokens_buffer.push(utils::fast_hash("https"));
             }
 
-            FilterTokens::Other(tokens.into_vec())
+            FilterTokens::Other(tokens_buffer.as_slice())
         }
     }
 }

src/flatbuffers/unsafe_tools.rs

@@ -101,7 +101,7 @@ impl VerifiedFlatbufferMemory {
 
 /// 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> {
+pub struct StackVector<T, const MAX_SIZE: usize> {
     data: [MaybeUninit<T>; MAX_SIZE],
     size: usize,
 }
@@ -133,6 +133,14 @@ impl<T, const MAX_SIZE: usize> StackVector<T, MAX_SIZE> {
         self.size == 0
     }
 
+    pub fn clear(&mut self) {
+        self.size = 0;
+    }
+
+    pub fn as_slice(&self) -> &[T] {
+        unsafe { std::slice::from_raw_parts(self.data.as_ptr() as *const T, self.size) }
+    }
+
     pub fn into_vec(self) -> Vec<T> {
         let mut v = Vec::with_capacity(self.size);
         for i in 0..self.size {

src/utils.rs

@@ -29,7 +29,7 @@ fn is_allowed_filter(ch: char) -> bool {
     ch.is_alphanumeric() || ch == '%'
 }
 
-pub(crate) type TokensBuffer = StackVector<Hash, 200>;
+pub type TokensBuffer = StackVector<Hash, 200>;
 
 fn fast_tokenizer_no_regex(
     pattern: &str,

tests/unit/filters/network.rs

@@ -1191,9 +1191,10 @@ mod parse_tests {
     fn test_simple_pattern_tokenization() {
         let rule = "||some.primewire.c*/sw$script,1p";
         let filter = NetworkFilter::parse(rule, true, ParseOptions::default()).unwrap();
+        let mut tokens_buffer = utils::TokensBuffer::default();
         assert_eq!(
-            filter.get_tokens_optimized(),
-            FilterTokens::Other(vec![
+            filter.get_tokens_optimized(&mut tokens_buffer),
+            FilterTokens::Other(&[
                 utils::fast_hash("some"),
                 utils::fast_hash("primewire")
             ])