Flat containers

src/data_format/mod.rs

@@ -10,7 +10,7 @@ mod storage;
 pub(crate) mod utils;
 
 use crate::cosmetic_filter_cache::CosmeticFilterCache;
-use crate::filters::unsafe_tools::VerifiedFlatbufferMemory;
+use crate::flatbuffers::unsafe_tools::VerifiedFlatbufferMemory;
 use crate::network_filter_list::NetworkFilterListParsingError;
 
 /// Newer formats start with this magic byte sequence.

src/data_format/storage.rs

@@ -10,7 +10,7 @@ use rmp_serde as rmps;
 use serde::{Deserialize, Serialize};
 
 use crate::cosmetic_filter_cache::{CosmeticFilterCache, HostnameRuleDb, ProceduralOrActionFilter};
-use crate::filters::unsafe_tools::VerifiedFlatbufferMemory;
+use crate::flatbuffers::unsafe_tools::VerifiedFlatbufferMemory;
 use crate::utils::Hash;
 
 use super::utils::{stabilize_hashmap_serialization, stabilize_hashset_serialization};

src/filters/fb_builder.rs

@@ -9,7 +9,7 @@ use std::vec;
 use flatbuffers::WIPOffset;
 
 use crate::filters::network::{NetworkFilter, NetworkFilterMaskHelper};
-use crate::filters::unsafe_tools::VerifiedFlatbufferMemory;
+use crate::flatbuffers::unsafe_tools::VerifiedFlatbufferMemory;
 use crate::network_filter_list::token_histogram;
 use crate::optimizer;
 use crate::utils::{to_short_hash, Hash, ShortHash};

src/filters/fb_network.rs

@@ -4,7 +4,7 @@ use std::collections::HashMap;
 
 use crate::filters::fb_builder::FlatBufferBuilder;
 use crate::filters::network::{NetworkFilterMask, NetworkFilterMaskHelper, NetworkMatchable};
-use crate::filters::unsafe_tools::{fb_vector_to_slice, VerifiedFlatbufferMemory};
+use crate::flatbuffers::unsafe_tools::{fb_vector_to_slice, VerifiedFlatbufferMemory};
 
 use crate::regex_manager::RegexManager;
 use crate::request::Request;

src/filters/mod.rs

@@ -6,6 +6,4 @@ mod network_matchers;
 pub mod cosmetic;
 pub(crate) mod fb_builder;
 pub(crate) mod fb_network;
-pub(crate) mod flat_filter_map;
 pub mod network;
-pub(crate) mod unsafe_tools;

src/flatbuffers/containers/flat_multimap.rs

@@ -0,0 +1,86 @@
+use std::marker::PhantomData;
+
+use crate::flatbuffers::containers::sorted_index::SortedIndex;
+use flatbuffers::{Follow, Vector};
+
+/// A map-like container that uses flatbuffer references.
+/// Provides O(log n) lookup time using binary search on the sorted index.
+/// I is a key type, Keys is specific container of keys, &[I] for fast indexing (u32, u64)
+/// and flatbuffers::Vector<I> if there is no conversion from Vector (str) to slice.
+pub(crate) struct FlatMultiMapView<'a, I: Ord, V, Keys>
+where
+    Keys: SortedIndex<I>,
+    V: Follow<'a>,
+{
+    keys: Keys,
+    values: Vector<'a, V>,
+    _phantom: PhantomData<I>,
+}
+
+impl<'a, I: Ord + Copy, V, Keys> FlatMultiMapView<'a, I, V, Keys>
+where
+    Keys: SortedIndex<I> + Clone,
+    V: Follow<'a>,
+{
+    pub fn new(keys: Keys, values: Vector<'a, V>) -> Self {
+        debug_assert!(keys.len() == values.len());
+
+        Self {
+            keys,
+            values,
+            _phantom: PhantomData,
+        }
+    }
+
+    pub fn get(&self, key: I) -> Option<FlatMultiMapViewIterator<'a, I, V, Keys>> {
+        let index = self.keys.partition_point(|x| *x < key);
+        if index < self.keys.len() && self.keys.get(index) == key {
+            Some(FlatMultiMapViewIterator {
+                index,
+                key,
+                keys: self.keys.clone(), // Cloning is 3-4% faster than & in benchmarks
+                values: self.values,
+            })
+        } else {
+            None
+        }
+    }
+
+    #[cfg(test)]
+    pub fn total_size(&self) -> usize {
+        self.keys.len()
+    }
+}
+
+pub(crate) struct FlatMultiMapViewIterator<'a, I: Ord + Copy, V, Keys>
+where
+    Keys: SortedIndex<I>,
+    V: Follow<'a>,
+{
+    index: usize,
+    key: I,
+    keys: Keys,
+    values: Vector<'a, V>,
+}
+
+impl<'a, I, V, Keys> Iterator for FlatMultiMapViewIterator<'a, I, V, Keys>
+where
+    I: Ord + Copy,
+    V: Follow<'a>,
+    Keys: SortedIndex<I>,
+{
+    type Item = (usize, <V as Follow<'a>>::Inner);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.index < self.keys.len() && self.keys.get(self.index) == self.key {
+            self.index += 1;
+            Some((self.index - 1, self.values.get(self.index - 1)))
+        } else {
+            None
+        }
+    }
+}
+
+#[cfg(test)]
+#[path = "../../../tests/unit/flatbuffers/containers/flat_multimap.rs"]
+mod unit_tests;

src/flatbuffers/containers/flat_set.rs

@@ -0,0 +1,49 @@
+#![allow(dead_code)]
+
+use std::marker::PhantomData;
+
+use crate::flatbuffers::containers::sorted_index::SortedIndex;
+
+/// A set-like container that uses flatbuffer references.
+/// Provides O(log n) lookup time using binary search on the sorted data.
+/// I is a key type, Keys is specific container of keys, &[I] for fast indexing (u32, u64)
+/// and flatbuffers::Vector<I> if there is no conversion from Vector (str) to slice.
+pub(crate) struct FlatSetView<I, Keys>
+where
+    Keys: SortedIndex<I>,
+{
+    keys: Keys,
+    _phantom: PhantomData<I>,
+}
+
+impl<I, Keys> FlatSetView<I, Keys>
+where
+    I: Ord,
+    Keys: SortedIndex<I>,
+{
+    pub fn new(keys: Keys) -> Self {
+        Self {
+            keys,
+            _phantom: PhantomData,
+        }
+    }
+
+    pub fn contains(&self, key: I) -> bool {
+        let index = self.keys.partition_point(|x| *x < key);
+        index < self.keys.len() && self.keys.get(index) == key
+    }
+
+    #[inline(always)]
+    pub fn len(&self) -> usize {
+        self.keys.len()
+    }
+
+    #[inline(always)]
+    pub fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+}
+
+#[cfg(test)]
+#[path = "../../../tests/unit/flatbuffers/containers/flat_set.rs"]
+mod unit_tests;

src/flatbuffers/containers/mod.rs

@@ -0,0 +1,3 @@
+pub(crate) mod flat_multimap;
+pub(crate) mod flat_set;
+pub(crate) mod sorted_index;

src/flatbuffers/containers/sorted_index.rs

@@ -0,0 +1,73 @@
+use flatbuffers::{Follow, Vector};
+
+// Represents sorted sequence to perform the binary search.
+pub(crate) trait SortedIndex<I> {
+    fn len(&self) -> usize;
+    fn get(&self, index: usize) -> I;
+    fn partition_point<F>(&self, predicate: F) -> usize
+    where
+        F: FnMut(&I) -> bool;
+}
+
+// Implementation for slices. Prefer using this with fb_vector_to_slice
+// if possible, because it faster than getting values with flatbuffer's
+// get method.
+impl<I: Ord + Copy> SortedIndex<I> for &[I] {
+    #[inline(always)]
+    fn len(&self) -> usize {
+        <[I]>::len(self)
+    }
+
+    #[inline(always)]
+    fn get(&self, index: usize) -> I {
+        self[index]
+    }
+
+    #[inline(always)]
+    fn partition_point<F>(&self, predicate: F) -> usize
+    where
+        F: FnMut(&I) -> bool,
+    {
+        debug_assert!(self.is_sorted());
+        <[I]>::partition_point(self, predicate)
+    }
+}
+
+// General implementation for flatbuffers::Vector, it uses get to
+// obtain values.
+impl<'a, T: Follow<'a>> SortedIndex<T::Inner> for Vector<'a, T>
+where
+    T::Inner: Ord,
+{
+    #[inline(always)]
+    fn len(&self) -> usize {
+        Vector::len(self)
+    }
+
+    #[inline(always)]
+    fn get(&self, index: usize) -> T::Inner {
+        Vector::get(self, index)
+    }
+
+    fn partition_point<F>(&self, mut predicate: F) -> usize
+    where
+        F: FnMut(&T::Inner) -> bool,
+    {
+        debug_assert!(self.iter().is_sorted());
+
+        let mut left = 0;
+        let mut right = self.len();
+
+        while left < right {
+            let mid = left + (right - left) / 2;
+            let value = self.get(mid);
+            if predicate(&value) {
+                left = mid + 1;
+            } else {
+                right = mid;
+            }
+        }
+
+        left
+    }
+}

src/flatbuffers/mod.rs

@@ -0,0 +1,2 @@
+pub(crate) mod containers;
+pub(crate) mod unsafe_tools;

src/lib.rs

@@ -23,6 +23,7 @@ pub mod cosmetic_filter_cache;
 mod data_format;
 mod engine;
 pub mod filters;
+mod flatbuffers;
 pub mod lists;
 mod network_filter_list;
 mod optimizer;