[perf] Replace token_histogram to token selector (#563)

atuchin-m · web-flow · commit 3feffc97ad17 · 2025-11-12T12:21:27.000+04:00
The PRs replaced old (slow and not enough memory preserving) token_histogram to TokenSelector. It
* process filters one by one, which faster and uses less memory. Also it unblocks the next optimizations.
* has the problematic tokens hardcoded;
* saves about -10% of building time.
* provides the final token quality that is a little worse, but the difference in matching performance is about 3-4%.
diff --git a/src/filters/fb_network_builder.rs b/src/filters/fb_network_builder.rs
@@ -6,11 +6,11 @@ use flatbuffers::WIPOffset;
 
 use crate::filters::fb_builder::EngineFlatBuilder;
 use crate::filters::network::{FilterTokens, NetworkFilter};
+use crate::filters::token_selector::TokenSelector;
 
 use crate::filters::network::NetworkFilterMaskHelper;
 use crate::flatbuffers::containers::flat_multimap::FlatMultiMapBuilder;
 use crate::flatbuffers::containers::flat_serialize::{FlatBuilder, FlatSerialize, WIPFlatVec};
-use crate::network_filter_list::token_histogram;
 use crate::optimizer;
 use crate::utils::{to_short_hash, Hash, ShortHash};
 
@@ -133,21 +133,10 @@ impl<'a> FlatSerialize<'a, EngineFlatBuilder<'a>> for NetworkFilterListBuilder {
 
         let mut optimizable = HashMap::<ShortHash, Vec<NetworkFilter>>::new();
 
-        // Compute tokens for all filters
-        let filter_tokens: Vec<_> = rule_list
-            .filters
-            .into_iter()
-            .map(|filter| {
-                let tokens = filter.get_tokens_optimized();
-                (filter, tokens)
-            })
-            .collect();
-
-        // compute the tokens' frequency histogram
-        let (total_number_of_tokens, tokens_histogram) = token_histogram(&filter_tokens);
+        let mut token_frequencies = TokenSelector::new(rule_list.filters.len());
 
         {
-            for (network_filter, multi_tokens) in filter_tokens.into_iter() {
+            for network_filter in rule_list.filters {
                 let flat_filter = if !rule_list.optimize
                     || !optimizer::is_filter_optimizable_by_patterns(&network_filter)
                 {
@@ -156,46 +145,34 @@ impl<'a> FlatSerialize<'a, EngineFlatBuilder<'a>> for NetworkFilterListBuilder {
                     None
                 };
 
-                let mut store_filter = |token: ShortHash| {
+                let mut store_filter = |token: Hash| {
+                    let short_token = to_short_hash(token);
                     if let Some(flat_filter) = flat_filter {
-                        filter_map_builder.insert(token, flat_filter);
+                        filter_map_builder.insert(short_token, flat_filter);
                     } else {
                         optimizable
-                            .entry(token)
+                            .entry(short_token)
                             .or_default()
                             .push(network_filter.clone());
                     }
                 };
 
+                let multi_tokens = network_filter.get_tokens_optimized();
                 match multi_tokens {
                     FilterTokens::Empty => {
-                        // No tokens, skip this filter
+                        // No tokens, add to fallback bucket (token 0)
+                        store_filter(0);
                     }
                     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));
+                            store_filter(token);
+                            token_frequencies.record_usage(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
-                                }
-                                _ => {}
-                            }
-                        }
-
+                        let best_token = token_frequencies.select_least_used_token(&tokens);
+                        token_frequencies.record_usage(best_token);
                         store_filter(best_token);
                     }
                 }
diff --git a/src/filters/mod.rs b/src/filters/mod.rs
@@ -9,6 +9,7 @@ pub(crate) mod fb_network;
 pub(crate) mod fb_network_builder;
 pub(crate) mod filter_data_context;
 pub mod network;
+pub(crate) mod token_selector;
 
 #[allow(unknown_lints)]
 #[allow(
diff --git a/src/filters/token_selector.rs b/src/filters/token_selector.rs
@@ -0,0 +1,117 @@
+//! Token selector for optimizing filter storage by choosing least-used tokens
+
+use crate::utils::{to_short_hash, Hash, ShortHash};
+use seahash::SeaHasher;
+use std::collections::HashMap;
+use std::hash::BuildHasherDefault;
+
+// Two groups of tokens that better be avoided during token selection.
+// The list was built by logging each token that triggered .matches() call
+// for some network filter in the benchmark.
+const WORST_TOKENS: [&str; 4] = ["https", "http", "www", "com"];
+const BAD_TOKENS: [&str; 36] = [
+    "uk",
+    "net",
+    "org",
+    "io",
+    "de",
+    "fr",
+    "es",
+    "it",
+    "nl",
+    "se",
+    "ru",
+    "pl",
+    "co",
+    "js",
+    "css",
+    "img",
+    "jpg",
+    "html",
+    "png",
+    "cdn",
+    "static",
+    "images",
+    "api",
+    "wp",
+    "ad",
+    "ads",
+    "content",
+    "doubleclick",
+    "analytics",
+    "assets",
+    "id",
+    "min",
+    "amazon",
+    "google",
+    "googlesyndication",
+    "googleapis",
+];
+
+const WORST_TOKEN_USAGE: usize = usize::MAX / 2;
+const BAD_TOKEN_USAGE: usize = usize::MAX / 4;
+
+/// Selects the optimal token for filter storage by tracking usage frequencies.
+/// Tokens that are used less frequently are preferred for better efficiency.
+pub(crate) struct TokenSelector {
+    usage: HashMap<ShortHash, usize, BuildHasherDefault<SeaHasher>>,
+}
+
+impl TokenSelector {
+    /// Creates a new TokenSelector with pre-populated commonly-used tokens.
+    pub fn new(capacity: usize) -> Self {
+        let mut usage = HashMap::with_capacity_and_hasher(
+            capacity + WORST_TOKENS.len() + BAD_TOKENS.len(),
+            BuildHasherDefault::<SeaHasher>::default(),
+        );
+        let mut store_token = |token: &str, count: usize| {
+            usage.insert(to_short_hash(crate::utils::fast_hash(token)), count);
+        };
+
+        for token in WORST_TOKENS {
+            store_token(token, WORST_TOKEN_USAGE);
+        }
+
+        for token in BAD_TOKENS {
+            store_token(token, BAD_TOKEN_USAGE);
+        }
+
+        Self { usage }
+    }
+
+    /// Selects the least-used token from the provided list of tokens.
+    /// Returns the token with the lowest current usage count.
+    pub fn select_least_used_token(&self, tokens: &[Hash]) -> Hash {
+        let mut best_token = 0;
+        let mut best_token_used = usize::MAX;
+        for &token in tokens {
+            if token == 0 {
+                // 0 is already used as a fallback token.
+                continue;
+            }
+
+            match self.usage.get(&to_short_hash(token)) {
+                Some(&count) => {
+                    if count < best_token_used {
+                        best_token_used = count;
+                        best_token = token;
+                    }
+                }
+                None => {
+                    // Token never seen before - this is optimal
+                    return token;
+                }
+            }
+        }
+        best_token
+    }
+
+    /// Records that a token has been used, incrementing its usage count.
+    pub fn record_usage(&mut self, token: Hash) {
+        *self.usage.entry(to_short_hash(token)).or_insert(0) += 1;
+    }
+}
+
+#[cfg(test)]
+#[path = "../../tests/unit/filters/token_selector.rs"]
+mod unit_tests;
diff --git a/src/network_filter_list.rs b/src/network_filter_list.rs
@@ -1,20 +1,20 @@
 //! Holds the implementation of [NetworkFilterList] and related functionality.
 
-use std::{collections::HashMap, collections::HashSet, fmt};
+use std::{collections::HashSet, fmt};
 
 use flatbuffers::ForwardsUOffset;
 
 use crate::filters::fb_network::FlatNetworkFilter;
 use crate::filters::filter_data_context::FilterDataContext;
 use crate::filters::flatbuffer_generated::fb;
 use crate::filters::network::{
-    FilterTokens, NetworkFilter, NetworkFilterMask, NetworkFilterMaskHelper, NetworkMatchable,
+    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, ShortHash};
+use crate::utils::{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.
@@ -151,27 +151,3 @@ impl NetworkFilterList<'_> {
         filters
     }
 }
-
-pub(crate) fn token_histogram<T>(
-    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() {
-        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 => {}
-        }
-    }
-
-    for bad_token in ["http", "https", "www", "com"].iter() {
-        tokens_histogram.insert(to_short_hash(fast_hash(bad_token)), number_of_tokens);
-    }
-
-    (number_of_tokens, tokens_histogram)
-}
diff --git a/tests/unit/engine.rs b/tests/unit/engine.rs
@@ -237,9 +237,9 @@ mod tests {
             );
         }
         let expected_hash: u64 = if cfg!(feature = "css-validation") {
-            15545091389304905433
+            13277824246832611772
         } else {
-            543362704487480180
+            12001568478200869587
         };
 
         assert_eq!(hash(&data), expected_hash, "{HASH_MISMATCH_MSG}");
diff --git a/tests/unit/filters/token_selector.rs b/tests/unit/filters/token_selector.rs
@@ -0,0 +1,46 @@
+#[cfg(test)]
+mod token_selector_tests {
+    use super::super::*;
+    use crate::utils::fast_hash;
+
+    #[test]
+    fn token_priority() {
+        let selector = TokenSelector::new(0);
+        let regular = fast_hash("rare_token");
+        let worst = fast_hash("https");
+        let bad = fast_hash("assets");
+
+        assert_eq!(selector.select_least_used_token(&[]), 0);
+        assert_eq!(selector.select_least_used_token(&[0, 0]), 0);
+
+        // a regular token is always better
+        assert_eq!(
+            selector.select_least_used_token(&[regular, worst, bad]),
+            regular
+        );
+        assert_eq!(
+            selector.select_least_used_token(&[worst, bad, regular]),
+            regular
+        );
+
+        // a bad token is always better than a worst token
+        assert_eq!(selector.select_least_used_token(&[worst, bad]), bad);
+        assert_eq!(selector.select_least_used_token(&[bad, worst]), bad);
+    }
+
+    #[test]
+    fn test_select_least_used_token_with_usage() {
+        let mut selector = TokenSelector::new(0);
+        let token1 = fast_hash("token1");
+        let token2 = fast_hash("token2");
+
+        assert_eq!(selector.select_least_used_token(&[token1, token2]), token1);
+
+        selector.record_usage(token1);
+        selector.record_usage(token1);
+        selector.record_usage(token2);
+
+        // token2 should be selected as it has lower usage
+        assert_eq!(selector.select_least_used_token(&[token1, token2]), token2);
+    }
+}
diff --git a/tests/unit/network_filter_list.rs b/tests/unit/network_filter_list.rs
@@ -26,49 +26,6 @@ mod tests {
         );
     }
 
-    #[test]
-    fn token_histogram_works() {
-        // handle the case of just 1 token
-        {
-            let tokens = vec![(0, vec![vec![111]])];
-            let (total_tokens, histogram) = token_histogram(&tokens);
-            assert_eq!(total_tokens, 1);
-            assert_eq!(histogram.get(&111), Some(&1));
-            // include bad tokens
-            assert_eq!(histogram.get(&fast_hash("http")), Some(&1));
-            assert_eq!(histogram.get(&fast_hash("www")), Some(&1));
-        }
-
-        // handle the case of repeating tokens
-        {
-            let tokens = vec![(0, vec![vec![111]]), (1, vec![vec![111]])];
-            let (total_tokens, histogram) = token_histogram(&tokens);
-            assert_eq!(total_tokens, 2);
-            assert_eq!(histogram.get(&111), Some(&2));
-            // include bad tokens
-            assert_eq!(histogram.get(&fast_hash("http")), Some(&2));
-            assert_eq!(histogram.get(&fast_hash("www")), Some(&2));
-        }
-
-        // handle the different token set sizes
-        {
-            let tokens = vec![
-                (0, vec![vec![111, 123, 132]]),
-                (1, vec![vec![111], vec![123], vec![132]]),
-                (2, vec![vec![111, 123], vec![132]]),
-                (3, vec![vec![111, 111], vec![111]]),
-            ];
-            let (total_tokens, histogram) = token_histogram(&tokens);
-            assert_eq!(total_tokens, 12);
-            assert_eq!(histogram.get(&111), Some(&6));
-            assert_eq!(histogram.get(&123), Some(&3));
-            assert_eq!(histogram.get(&132), Some(&3));
-            // include bad tokens
-            assert_eq!(histogram.get(&fast_hash("http")), Some(&12));
-            assert_eq!(histogram.get(&fast_hash("www")), Some(&12));
-        }
-    }
-
     #[test]
     fn network_filter_list_new_works() {
         {
