Speed up cosmetic filtering

Cargo.toml

@@ -34,6 +34,8 @@ idna = "1.0.3"
 serde = { workspace = true }
 serde_json = { workspace = true }
 seahash = "4.1.0"
+# rustc-hash v1.1.0 provides a better performance than 2.x, chromium pins the same version.
+rustc-hash = { version = "1.1.0", default-features = false }
 memchr = "2.4"
 base64 = "0.22"
 rmp-serde = "0.15"

src/cosmetic_filter_cache.rs

@@ -16,8 +16,9 @@ use crate::filters::cosmetic::{CosmeticFilterAction, CosmeticFilterOperator};
 use crate::filters::filter_data_context::FilterDataContextRef;
 
 use crate::flatbuffers::containers::flat_map::FlatMapView;
-use crate::flatbuffers::containers::flat_multimap::{FlatMapStringView, FlatMultiMapView};
-use crate::flatbuffers::containers::flat_set::FlatSetView;
+use crate::flatbuffers::containers::flat_multimap::FlatMultiMapView;
+use crate::flatbuffers::containers::hash_map::HashMapStringView;
+use crate::flatbuffers::containers::hash_set::HashSetView;
 use crate::resources::{PermissionMask, ResourceStorage};
 
 use crate::utils::Hash;
@@ -169,13 +170,13 @@ impl CosmeticFilterCache {
         let mut selectors = vec![];
 
         let cosmetic_filters = self.filter_data_context.memory.root().cosmetic_filters();
-        let simple_class_rules = FlatSetView::new(cosmetic_filters.simple_class_rules());
-        let simple_id_rules = FlatSetView::new(cosmetic_filters.simple_id_rules());
-        let complex_class_rules = FlatMapStringView::new(
+        let simple_class_rules = HashSetView::new(cosmetic_filters.simple_class_rules());
+        let simple_id_rules = HashSetView::new(cosmetic_filters.simple_id_rules());
+        let complex_class_rules = HashMapStringView::new(
             cosmetic_filters.complex_class_rules_index(),
             cosmetic_filters.complex_class_rules_values(),
         );
-        let complex_id_rules = FlatMapStringView::new(
+        let complex_id_rules = HashMapStringView::new(
             cosmetic_filters.complex_id_rules_index(),
             cosmetic_filters.complex_id_rules_values(),
         );
@@ -185,8 +186,8 @@ impl CosmeticFilterCache {
             if simple_class_rules.contains(class) && !exceptions.contains(&format!(".{}", class)) {
                 selectors.push(format!(".{}", class));
             }
-            if let Some(bucket) = complex_class_rules.get(class) {
-                for (_, sel) in bucket {
+            if let Some(values) = complex_class_rules.get(class) {
+                for sel in values.data() {
                     if !exceptions.contains(sel) {
                         selectors.push(sel.to_string());
                     }
@@ -198,8 +199,8 @@ impl CosmeticFilterCache {
             if simple_id_rules.contains(id) && !exceptions.contains(&format!("#{}", id)) {
                 selectors.push(format!("#{}", id));
             }
-            if let Some(bucket) = complex_id_rules.get(id) {
-                for (_, sel) in bucket {
+            if let Some(values) = complex_id_rules.get(id) {
+                for sel in values.data() {
                     if !exceptions.contains(sel) {
                         selectors.push(sel.to_string());
                     }

src/cosmetic_filter_cache_builder.rs

@@ -9,6 +9,8 @@ use crate::filters::cosmetic::{CosmeticFilter, CosmeticFilterMask, CosmeticFilte
 use crate::filters::flatbuffer_generated::fb;
 use crate::flatbuffers::containers::flat_map::FlatMapBuilder;
 use crate::flatbuffers::containers::flat_multimap::FlatMultiMapBuilder;
+use crate::flatbuffers::containers::hash_map::HashMapBuilder;
+use crate::flatbuffers::containers::hash_set::HashSetBuilder;
 
 use crate::flatbuffers::containers::flat_serialize::{
     serialize_vec_opt, FlatBuilder, FlatSerialize,
@@ -56,13 +58,16 @@ impl<'a, B: FlatBuilder<'a>> FlatSerialize<'a, B> for HostnameRule {
     }
 }
 
+#[derive(Default, Clone)]
+struct StringVector(Vec<String>);
+
 #[derive(Default)]
 pub(crate) struct CosmeticFilterCacheBuilder {
-    simple_class_rules: HashSet<String>,
-    simple_id_rules: HashSet<String>,
+    simple_class_rules: HashSetBuilder<String>,
+    simple_id_rules: HashSetBuilder<String>,
     misc_generic_selectors: HashSet<String>,
-    complex_class_rules: FlatMultiMapBuilder<String, String>,
-    complex_id_rules: FlatMultiMapBuilder<String, String>,
+    complex_class_rules: HashMapBuilder<String, StringVector>,
+    complex_id_rules: HashMapBuilder<String, StringVector>,
 
     hostname_hide: FlatMultiMapBuilder<Hash, String>,
     hostname_inject_script: FlatMultiMapBuilder<Hash, String>,
@@ -110,7 +115,10 @@ impl CosmeticFilterCacheBuilder {
                 if key == selector {
                     self.simple_class_rules.insert(class);
                 } else {
-                    self.complex_class_rules.insert(class, selector);
+                    let selectors = self
+                        .complex_class_rules
+                        .get_or_insert(class, StringVector::default());
+                    selectors.0.push(selector);
                 }
             }
         } else if selector.starts_with('#') {
@@ -120,7 +128,10 @@ impl CosmeticFilterCacheBuilder {
                 if key == selector {
                     self.simple_id_rules.insert(id);
                 } else {
-                    self.complex_id_rules.insert(id, selector);
+                    let selectors = self
+                        .complex_id_rules
+                        .get_or_insert(id, StringVector::default());
+                    selectors.0.push(selector);
                 }
             }
         } else {
@@ -204,11 +215,24 @@ impl CosmeticFilterCacheBuilder {
     }
 }
 
+impl<'a, B: FlatBuilder<'a>> FlatSerialize<'a, B> for StringVector {
+    type Output = WIPOffset<fb::StringVector<'a>>;
+
+    fn serialize(value: Self, builder: &mut B) -> WIPOffset<fb::StringVector<'a>> {
+        let v = FlatSerialize::serialize(value.0, builder);
+        fb::StringVector::create(
+            builder.raw_builder(),
+            &fb::StringVectorArgs { data: Some(v) },
+        )
+    }
+}
+
 impl<'a, B: FlatBuilder<'a>> FlatSerialize<'a, B> for CosmeticFilterCacheBuilder {
     type Output = WIPOffset<fb::CosmeticFilters<'a>>;
+
     fn serialize(value: Self, builder: &mut B) -> WIPOffset<fb::CosmeticFilters<'a>> {
-        let complex_class_rules = FlatMultiMapBuilder::finish(value.complex_class_rules, builder);
-        let complex_id_rules = FlatMultiMapBuilder::finish(value.complex_id_rules, builder);
+        let complex_class_rules = HashMapBuilder::finish(value.complex_class_rules, builder);
+        let complex_id_rules = HashMapBuilder::finish(value.complex_id_rules, builder);
 
         // Handle top-level hostname hide and inject_script for better deduplication
         let hostname_hide = FlatMultiMapBuilder::finish(value.hostname_hide, builder);

src/data_format/mod.rs

@@ -17,7 +17,7 @@ const ADBLOCK_RUST_DAT_MAGIC: [u8; 4] = [0xd1, 0xd9, 0x3a, 0xaf];
 
 /// The version of the data format.
 /// If the data format version is incremented, the data is considered as incompatible.
-const ADBLOCK_FLATBUFFER_VERSION: u8 = 2;
+const ADBLOCK_RUST_DAT_VERSION: u8 = 2;
 
 /// The total length of the header prefix (magic + version + seahash)
 const HEADER_PREFIX_LENGTH: usize = 4 + 1 + 8;
@@ -35,7 +35,7 @@ pub(crate) fn serialize_dat_file(data: &[u8]) -> Vec<u8> {
     let mut serialized = Vec::with_capacity(data.len() + HEADER_PREFIX_LENGTH);
     let hash = seahash::hash(data).to_le_bytes();
     serialized.extend_from_slice(&ADBLOCK_RUST_DAT_MAGIC);
-    serialized.push(ADBLOCK_FLATBUFFER_VERSION);
+    serialized.push(ADBLOCK_RUST_DAT_VERSION);
     serialized.extend_from_slice(&hash);
     assert_eq!(serialized.len(), HEADER_PREFIX_LENGTH);
 
@@ -49,7 +49,7 @@ pub(crate) fn deserialize_dat_file(serialized: &[u8]) -> Result<&[u8], Deseriali
     }
 
     let version = serialized[ADBLOCK_RUST_DAT_MAGIC.len()];
-    if version != ADBLOCK_FLATBUFFER_VERSION {
+    if version != ADBLOCK_RUST_DAT_VERSION {
         return Err(DeserializationError::VersionMismatch(version));
     }
     let data = &serialized[HEADER_PREFIX_LENGTH..];

src/flatbuffers/containers/fb_index.rs

@@ -0,0 +1,49 @@
+use flatbuffers::{Follow, Vector};
+
+/// A trait to access indexed data in a flatbuffer.
+/// It has two implementations:
+/// 1. a faster &[I] for slices;
+/// 2. a slower for flatbuffers::Vector<I>, that uses Follow() internally.
+///
+/// Note: it intentionally returns values using a copy, because it's faster
+/// than by reference.
+pub(crate) trait FbIndex<I> {
+    /// Returns the number of elements.
+    fn len(&self) -> usize;
+
+    /// Returns a copy of the value at the given index.
+    /// 'index' must be in range [0, len()), otherwise panics.
+    fn get(&self, index: usize) -> I;
+}
+
+impl<I: Copy> FbIndex<I> for &[I] {
+    #[inline(always)]
+    fn len(&self) -> usize {
+        <[I]>::len(self)
+    }
+
+    #[inline(always)]
+    fn get(&self, index: usize) -> I {
+        self[index]
+    }
+}
+
+impl FbIndex<()> for () {
+    #[inline(always)]
+    fn len(&self) -> usize {
+        0
+    }
+    fn get(&self, _index: usize) {}
+}
+
+impl<'a, T: Follow<'a>> FbIndex<T::Inner> for Vector<'a, T> {
+    #[inline(always)]
+    fn len(&self) -> usize {
+        Vector::len(self)
+    }
+
+    #[inline(always)]
+    fn get(&self, index: usize) -> T::Inner {
+        Vector::get(self, index)
+    }
+}

src/flatbuffers/containers/flat_multimap.rs

@@ -129,9 +129,6 @@ impl<I: Ord + std::hash::Hash, V> FlatMultiMapBuilder<I, V> {
     }
 }
 
-pub(crate) type FlatMapStringView<'a, V> =
-    FlatMultiMapView<'a, &'a str, V, Vector<'a, flatbuffers::ForwardsUOffset<&'a str>>>;
-
 #[cfg(test)]
 #[path = "../../../tests/unit/flatbuffers/containers/flat_multimap.rs"]
 mod unit_tests;

src/flatbuffers/containers/flat_serialize.rs

@@ -27,8 +27,13 @@ pub trait FlatSerialize<'b, B: FlatBuilder<'b>>: Sized {
 
 impl<'b> FlatBuilder<'b> for flatbuffers::FlatBufferBuilder<'b> {
     fn create_string(&mut self, s: &str) -> WIPOffset<&'b str> {
-        self.create_string(s)
+        if s.is_empty() {
+            flatbuffers::FlatBufferBuilder::create_shared_string(self, s)
+        } else {
+            flatbuffers::FlatBufferBuilder::create_string(self, s)
+        }
     }
+
     fn raw_builder(&mut self) -> &mut flatbuffers::FlatBufferBuilder<'b> {
         self
     }