Optimize network filter get tokens (#540)

* optimize NetworkFilter::get_tokens()
* Fix some review issues
* Return get_tokens() for API compatibility

Author: atuchin-m

src/filters/fb_network_builder.rs

@@ -5,7 +5,7 @@ use std::collections::{HashMap, HashSet};
 use flatbuffers::WIPOffset;
 
 use crate::filters::fb_builder::EngineFlatBuilder;
-use crate::filters::network::NetworkFilter;
+use crate::filters::network::{FilterTokens, NetworkFilter};
 
 use crate::filters::network::NetworkFilterMaskHelper;
 use crate::flatbuffers::containers::flat_multimap::FlatMultiMapBuilder;
@@ -138,7 +138,7 @@ impl<'a> FlatSerialize<'a, EngineFlatBuilder<'a>> for NetworkFilterListBuilder {
             .filters
             .into_iter()
             .map(|filter| {
-                let tokens = filter.get_tokens();
+                let tokens = filter.get_tokens_optimized();
                 (filter, tokens)
             })
             .collect();
@@ -156,33 +156,49 @@ impl<'a> FlatSerialize<'a, EngineFlatBuilder<'a>> for NetworkFilterListBuilder {
                     None
                 };
 
-                for tokens in multi_tokens {
-                    let mut best_token: ShortHash = 0;
-                    let mut min_count = total_number_of_tokens + 1;
-                    for token in tokens {
-                        let token = to_short_hash(token);
-                        match tokens_histogram.get(&token) {
-                            None => {
-                                min_count = 0;
-                                best_token = token
-                            }
-                            Some(&count) if count < min_count => {
-                                min_count = count;
-                                best_token = token
-                            }
-                            _ => {}
-                        }
-                    }
-
+                let mut store_filter = |token: ShortHash| {
                     if let Some(flat_filter) = flat_filter {
-                        filter_map.entry(best_token).or_default().push(flat_filter);
+                        filter_map.entry(token).or_default().push(flat_filter);
                     } else {
                         optimizable
-                            .entry(best_token)
+                            .entry(token)
                             .or_default()
                             .push(network_filter.clone());
                     }
-                } // tokens
+                };
+
+                match multi_tokens {
+                    FilterTokens::Empty => {
+                        // No tokens, skip this filter
+                    }
+                    FilterTokens::OptDomains(opt_domains) => {
+                        // For OptDomains, each domain is treated as a separate token group
+                        for &token in opt_domains.iter() {
+                            store_filter(to_short_hash(token));
+                        }
+                    }
+                    FilterTokens::Other(tokens) => {
+                        // For Other tokens, find the best token from the group
+                        let mut best_token: ShortHash = 0;
+                        let mut min_count = total_number_of_tokens + 1;
+                        for &token in tokens.iter() {
+                            let token = to_short_hash(token);
+                            match tokens_histogram.get(&token) {
+                                None => {
+                                    min_count = 0;
+                                    best_token = token
+                                }
+                                Some(&count) if count < min_count => {
+                                    min_count = count;
+                                    best_token = token
+                                }
+                                _ => {}
+                            }
+                        }
+
+                        store_filter(best_token);
+                    }
+                }
             }
         }
 

src/filters/network.rs

@@ -311,6 +311,12 @@ pub enum FilterPart {
     AnyOf(Vec<String>),
 }
 
+pub enum FilterTokens {
+    Empty,
+    OptDomains(Vec<Hash>),
+    Other(Vec<Hash>),
+}
+
 pub struct FilterPartIterator<'a> {
     filter_part: &'a FilterPart,
     index: usize,
@@ -876,7 +882,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() {
+            FilterTokens::OptDomains(domains) => {
+                domains.into_iter().map(|domain| vec![domain]).collect()
+            }
+            FilterTokens::Other(tokens) => vec![tokens],
+            FilterTokens::Empty => vec![],
+        }
+    }
+
+    pub fn get_tokens_optimized(&self) -> FilterTokens {
         let mut tokens: Vec<Hash> = Vec::with_capacity(TOKENS_BUFFER_SIZE);
 
         // If there is only one domain and no domain negation, we also use this
@@ -930,21 +947,24 @@ impl NetworkFilter {
         // 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() {
-            self.opt_domains
-                .as_ref()
-                .unwrap_or(&vec![])
-                .iter()
-                .map(|&d| vec![d])
-                .collect()
+            if let Some(opt_domains) = self.opt_domains.as_ref() {
+                if !opt_domains.is_empty() {
+                    return FilterTokens::OptDomains(opt_domains.clone());
+                }
+            }
+            FilterTokens::Empty
         } else {
             // Add optional token for protocol
             if self.for_http() && !self.for_https() {
                 tokens.push(utils::fast_hash("http"));
             } else if self.for_https() && !self.for_http() {
                 tokens.push(utils::fast_hash("https"));
             }
-            tokens.shrink_to_fit();
-            vec![tokens]
+
+            // Remake a vector to drop extra capacity.
+            let mut t = Vec::with_capacity(tokens.len());
+            t.extend(tokens);
+            FilterTokens::Other(t)
         }
     }
 }

src/network_filter_list.rs

@@ -8,13 +8,13 @@ use crate::filters::fb_network::FlatNetworkFilter;
 use crate::filters::filter_data_context::FilterDataContext;
 use crate::filters::flatbuffer_generated::fb;
 use crate::filters::network::{
-    NetworkFilter, NetworkFilterMask, NetworkFilterMaskHelper, NetworkMatchable,
+    FilterTokens, NetworkFilter, NetworkFilterMask, NetworkFilterMaskHelper, NetworkMatchable,
 };
 use crate::flatbuffers::containers::flat_multimap::FlatMultiMapView;
 use crate::flatbuffers::unsafe_tools::fb_vector_to_slice;
 use crate::regex_manager::RegexManager;
 use crate::request::Request;
-use crate::utils::{fast_hash, to_short_hash, Hash, ShortHash};
+use crate::utils::{fast_hash, to_short_hash, ShortHash};
 
 /// Holds relevant information from a single matchin gnetwork filter rule as a result of querying a
 /// [NetworkFilterList] for a given request.
@@ -155,16 +155,19 @@ impl NetworkFilterList<'_> {
 }
 
 pub(crate) fn token_histogram<T>(
-    filter_tokens: &[(T, Vec<Vec<Hash>>)],
+    filter_tokens: &[(T, FilterTokens)],
 ) -> (u32, HashMap<ShortHash, u32>) {
     let mut tokens_histogram: HashMap<ShortHash, u32> = HashMap::new();
     let mut number_of_tokens = 0;
     for (_, tokens) in filter_tokens.iter() {
-        for tg in tokens {
-            for t in tg {
-                *tokens_histogram.entry(to_short_hash(*t)).or_insert(0) += 1;
-                number_of_tokens += 1;
+        match tokens {
+            FilterTokens::Other(tokens) | FilterTokens::OptDomains(tokens) => {
+                for t in tokens {
+                    *tokens_histogram.entry(to_short_hash(*t)).or_insert(0) += 1;
+                    number_of_tokens += 1;
+                }
             }
+            FilterTokens::Empty => {}
         }
     }