Revision control
Copy as Markdown
Other Tools
/* This Source Code Form is subject to the terms of the Mozilla Public
* License, v. 2.0. If a copy of the MPL was not distributed with this
use crate::{
clubcard::ClubcardIndex, Clubcard, ClubcardIndexEntry, Equation, Filterable, Queryable,
};
use rand::{thread_rng, Rng};
use std::collections::BTreeMap;
use std::fmt;
/// Marker type for checking that, for example, only Exact ribbons are passed to functions such as
/// Clubcard::collect_exact_ribbons.
pub struct Exact;
/// A Ribbon Filter that encodes a one bit value for every element of the associated universe.
pub type ExactRibbon<const W: usize, T> = Ribbon<W, T, Exact>;
/// Marker type for checking that, for example, only Approximate ribbons are passed to functions such as
/// Clubcard::collect_approximate_ribbons.
pub struct Approximate;
/// A Ribbon Filter that identifies a subset of a universe with a false positive rate of
/// roughly |subset| / |universe|.
pub type ApproximateRibbon<const W: usize, T> = Ribbon<W, T, Approximate>;
/// A RibbonBuilder collects a set of items for insertion into a Ribbon. If the optional filter is
/// provided, then only items that are contained in the filter will be inserted.
pub struct RibbonBuilder<'a, const W: usize, T: Filterable<W>> {
/// block id.
id: Vec<u8>,
/// items to be inserted.
items: Vec<T>,
/// filter for pruning insertions.
filter: Option<&'a PartitionedRibbonFilter<W, T, Approximate>>,
/// size of the universe that contains self.items
universe_size: usize,
/// Whether queries against this ribbon indicate membership in R (inverted = false) or
/// membership in U \ R (inverted = true).
inverted: bool,
}
impl<'a, const W: usize, T: Filterable<W>> RibbonBuilder<'a, W, T> {
fn new(
id: &[u8],
filter: Option<&'a PartitionedRibbonFilter<W, T, Approximate>>,
) -> RibbonBuilder<'a, W, T> {
RibbonBuilder {
id: AsRef::<[u8]>::as_ref(id).to_vec(),
items: vec![],
filter,
universe_size: 0,
inverted: false,
}
}
/// Queue `item` for insertion into the ribbon (if it is contained in the provided filter).
pub fn insert(&mut self, item: T) {
if let Some(filter) = self.filter {
if filter.contains(&item) {
self.items.push(item);
}
} else {
self.items.push(item);
}
}
/// Set the size of the universe. This only needs to be called if you
/// are constructing an ApproximateRibbon.
pub fn set_universe_size(&mut self, universe_size: usize) {
self.universe_size = universe_size;
}
}
impl<'a, const W: usize, T: Filterable<W>> From<RibbonBuilder<'a, W, T>>
for ApproximateRibbon<W, T>
{
/// Denote the inserted set by R and the universe by U.
/// The ribbon returned by ApproximateRibbon::from encodes a function f : U -> {0, 1} where
/// f(x) = 0 if and only if x is in R union S where S is a (random) subset of U \ R of size
/// ~|R|. In other words, the ribbon solves the approximate membership query problem with a
/// false positive rate roughly 2^-r = |R| / (|U| - |R|).
/// The size of this ribbon is proportional to r|R|.
fn from(mut builder: RibbonBuilder<'a, W, T>) -> ApproximateRibbon<W, T> {
assert!(builder.items.len() <= builder.universe_size);
if builder.items.len() == builder.universe_size {
ApproximateRibbon::new(&builder.id, 0, builder.universe_size, !builder.inverted)
} else {
let mut out = ApproximateRibbon::new(
&builder.id,
builder.items.len(),
builder.universe_size,
builder.inverted,
);
for item in builder.items.drain(..) {
out.insert(item);
}
// Insertions should not fail for a homogeneous system.
assert!(out.exceptions.is_empty());
out
}
}
}
impl<'a, const W: usize, T: Filterable<W>> From<RibbonBuilder<'a, W, T>> for ExactRibbon<W, T> {
/// Denote the inserted set by R and the universe by U.
/// The ribbon returned by ExactRibbon::from encodes the function "f(x) = 0 iff x in R". The
/// size of this ribbon is proportional to |U|. In the typical use case, the set U is the
/// result of filtering a larger universe with a false positive rate of 2^-r. This allows for
/// exact encoding of R-membership using a pair of filters of total size ~(r+2)|R|.
fn from(mut builder: RibbonBuilder<'a, W, T>) -> ExactRibbon<W, T> {
assert!(builder.universe_size == 0 || builder.universe_size == builder.items.len());
if let Some(filter) = builder.filter {
if filter.block_is_empty(&builder.id) {
// The approximate filter is empty, so it gives a definitive result on every
// item and there's nothing to encode in the exact filter.
return ExactRibbon::new(&builder.id, 0, filter.block_is_inverted(&builder.id));
}
}
let mut out = ExactRibbon::new(&builder.id, builder.items.len(), builder.inverted);
// By inserting the included items first, we ensure that any exceptions that occur during
// insertion are for excluded items.
let mut excluded = vec![];
for item in builder.items.drain(..) {
if item.included() {
out.insert(item);
} else {
excluded.push(item);
}
}
for item in excluded.drain(..) {
out.insert(item);
}
out
}
}
/// A compact representation of a linear system AX = B
pub struct Ribbon<const W: usize, T: Filterable<W>, ApproxOrExact> {
/// A block identifier. Used to build an index for partitioned filters.
id: Vec<u8>,
/// The overhead.
epsilon: f64,
/// Equal to (1+epsilon) * |R|
m: usize,
/// The rank is round(-log2(subset_size / (universe_size - subset_size)))
rank: usize,
/// A linear system in which each equation has s in {0, ..., m-1}
rows: Vec<Equation<W>>,
/// A (typically short) list of items that failed insertion
exceptions: Vec<T>,
/// Whether queries against this ribbon indicate membership in R (inverted = false) or
/// membership in U \ R (inverted = true).
inverted: bool,
/// Marker for whether this is an Approximate or an Exact filter.
phantom: std::marker::PhantomData<ApproxOrExact>,
}
impl<const W: usize, T: Filterable<W>, ApproxOrExact> fmt::Display for Ribbon<W, T, ApproxOrExact> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(
f,
"ribbon({:?}): m: {}, rows: {}, rank: {}, exceptions: {}, epsilon: {}, overhead {}",
self.id,
self.m,
self.rows.len(),
self.rank,
self.exceptions.len(),
self.epsilon,
(self.rows.iter().filter(|eq| eq.is_zero()).count() as f64 / (self.rows.len() as f64))
)
}
}
impl<const W: usize, T: Filterable<W>> ApproximateRibbon<W, T> {
/// Construct an empty ribbon to encode a set R of size `subset_size` in a universe U of size
/// `universe_size`.
fn new(id: &[u8], subset_size: usize, universe_size: usize, inverted: bool) -> Self {
assert!(subset_size <= universe_size);
// TODO: Tune epsilon as a function of the inputs. Numerical experiments?
let epsilon = 0.02;
let m = ((1.0 + epsilon) * (subset_size as f64)).floor() as usize;
let rank = if subset_size == 0 || 2 * subset_size >= universe_size {
0
} else {
(((universe_size - subset_size) as f64) / (subset_size as f64))
.log2()
.floor() as usize
};
Ribbon {
id: AsRef::<[u8]>::as_ref(id).to_vec(),
rows: vec![Equation::zero(); m],
m,
epsilon,
rank,
exceptions: vec![],
inverted,
phantom: std::marker::PhantomData,
}
}
}
impl<const W: usize, T: Filterable<W>> ExactRibbon<W, T> {
/// Construct an empty ribbon to encode a set R of size `subset_size` in a universe U of size
/// `universe_size`.
fn new(id: &impl AsRef<[u8]>, size: usize, inverted: bool) -> Self {
// TODO: Tune epsilon as a function of the inputs. Numerical experiments?
let epsilon = 0.02;
let m = ((1.0 + epsilon) * (size as f64)).floor() as usize;
Ribbon {
id: AsRef::<[u8]>::as_ref(id).to_vec(),
rows: vec![Equation::zero(); m],
m,
epsilon,
rank: 1,
exceptions: vec![],
inverted,
phantom: std::marker::PhantomData,
}
}
}
impl<const W: usize, T: Filterable<W>, ApproxOrExact> Ribbon<W, T, ApproxOrExact> {
/// Hash the item to an Equation and insert it into the system.
fn insert(&mut self, item: T) -> bool {
let mut eq = item.as_query(self.m);
eq.b = if item.included() { 0 } else { 1 };
assert!(eq.is_zero() || eq.a[0] & 1 == 1);
let rv = self.insert_equation(eq);
if !rv {
self.exceptions.push(item)
}
rv
}
fn insert_equation(&mut self, mut eq: Equation<W>) -> bool {
loop {
if eq.is_zero() {
return eq.b == 0; /* redundant (b=0) or inconsistent (b!=0) */
}
if eq.s >= self.rows.len() {
// TODO: could be smarter here
self.rows.resize_with(eq.s + 1, Equation::zero);
}
let cur = &mut self.rows[eq.s];
if cur.is_zero() {
*cur = eq;
return true; /* inserted */
}
eq.add(cur);
}
}
/// Solve the system using back-substitution. If this is a block in a larger system, the `tail`
/// argument should be set to the the solution vector for the block to the right of this one.
fn solve(&self, tail: &[u64]) -> Vec<u64> {
let mut z = vec![0u64; ((self.rows.len() + 63) / 64) + tail.len()];
// insert tail into z starting at bit self.rows.len()
let k = self.rows.len() / 64;
let p = self.rows.len() % 64;
if p == 0 {
z[k..(tail.len() + k)].copy_from_slice(tail);
} else {
for i in 0..tail.len() {
z[k + i] |= tail[i] << p;
z[k + i + 1] = tail[i] >> (64 - p)
}
}
// Solve by back substitution
for i in (0..self.rows.len()).rev() {
let limb = i / 64;
let pos = i % 64;
let z_i = if self.rows[i].is_zero() {
// Row i has a zero in column i, so we're free to choose.
// TODO: We want multiple calls to solve() to give a different
// solutions (when the system is suitably under-determined),
// but it might be nice if this was deterministic.
thread_rng().gen::<u8>()
} else {
// Let z' be the vector we get by setting bit i of z to z'_i.
// Since z_i is zero, and row i has a one in column i, we have
// row_i(z') = z'_i ^ row_i(z).
// We want row_i(z') = b, so we must choose
// z'_i = row_i(z) ^ b.
self.rows[i].eval(&z) ^ self.rows[i].b
};
z[limb] |= ((z_i & 1) as u64) << pos;
}
z
}
}
#[derive(Debug)]
struct PartitionedRibbonFilterIndexEntry {
offset: usize,
m: usize,
rank: usize,
exceptions: Vec<Vec<u8>>,
inverted: bool,
}
type PartitionedRibbonFilterIndex =
BTreeMap</* block id */ Vec<u8>, PartitionedRibbonFilterIndexEntry>;
/// A solution to a ribbon system, along with metadata necessary for querying it.
struct PartitionedRibbonFilter<const W: usize, T: Filterable<W>, ApproxOrExact> {
index: PartitionedRibbonFilterIndex,
solution: Vec<Vec<u64>>,
phantom: std::marker::PhantomData<T>,
phantom2: std::marker::PhantomData<ApproxOrExact>,
}
impl<const W: usize, T: Filterable<W>, ApproxOrExact> fmt::Display
for PartitionedRibbonFilter<W, T, ApproxOrExact>
{
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "PartitionedRibbonFilter({:?})", self.index)
}
}
impl<const W: usize, T: Filterable<W>, Approximate> PartitionedRibbonFilter<W, T, Approximate> {
fn block_is_empty(&self, block: &[u8]) -> bool {
let Some(entry) = self.index.get(block) else {
return false;
};
entry.m == 0
}
fn block_is_inverted(&self, block: &[u8]) -> bool {
let Some(entry) = self.index.get(block) else {
return false;
};
entry.inverted
}
/// Check if this filter contains the given item in the given block.
fn contains(&self, item: &T) -> bool {
let Some(entry) = self.index.get(item.block()) else {
return false;
};
let result = (|| {
// Empty blocks do not contain anything,
// despite having inner product 0 with everything.
if entry.m == 0 {
return false;
}
let mut eq = item.as_query(entry.m);
eq.s += entry.offset;
for i in 0..entry.rank {
if eq.eval(&self.solution[i]) != 0 {
return false;
}
}
for exception in &entry.exceptions {
if exception == item.discriminant() {
return false;
}
}
true
})();
result ^ entry.inverted
}
}
impl<const W: usize, T: Filterable<W>, ApproxOrExact> From<Vec<Ribbon<W, T, ApproxOrExact>>>
for PartitionedRibbonFilter<W, T, ApproxOrExact>
{
fn from(
mut blocks: Vec<Ribbon<W, T, ApproxOrExact>>,
) -> PartitionedRibbonFilter<W, T, ApproxOrExact> {
// Sort ribbons by descending rank (descending simplifies indexing).
blocks.sort_unstable_by(|a, b| b.rank.cmp(&a.rank));
// Solve the (block) system.
// The blocks are sorted by descending rank. We need at least one solution (i.e. column
// vector) per block, but we need no more than i solutions for a block of rank i.
// Concretely, suppose the ranks are [4, 2, 1, 0]. Then our solution can look like
// block 0: | | | |
// block 1: | | 0 0
// block 2: | 0 0 0
// block 3: | 0 0 0
// Since we serialize the block identifiers, offsets, and ranks in the final filter, we
// don't need to encode the zeros.
let mut solution = vec![];
let max_rank = blocks.first().map_or(0, |first| first.rank);
for i in 0..max_rank {
// Back substitution across blocks.
let mut tail = vec![];
if max_rank > 1 {
// randomizing the tail increases the odds that the solutions will be distinct
tail.push(thread_rng().gen::<u64>());
}
for j in (0..blocks.len()).rev() {
if blocks[j].rank > i {
tail = blocks[j].solve(&tail);
}
}
while let Some(0) = tail.last() {
tail.pop();
}
solution.push(tail);
}
// construct the index---a map from a block identifier to that
// block's offset in the solution vector.
let mut index = PartitionedRibbonFilterIndex::new();
let mut offset = 0;
for block in &blocks {
let exceptions = block
.exceptions
.iter()
.map(|x| x.discriminant().to_vec())
.collect();
index.insert(
block.id.clone(),
PartitionedRibbonFilterIndexEntry {
offset,
m: block.m,
rank: block.rank,
exceptions,
inverted: block.inverted,
},
);
offset += block.rows.len();
}
PartitionedRibbonFilter {
index,
solution,
phantom: std::marker::PhantomData,
phantom2: std::marker::PhantomData,
}
}
}
/// A pair of ribbon filters that, together, solve the exact membership query problem.
pub struct ClubcardBuilder<const W: usize, T: Filterable<W>> {
/// An approximate membership query filter to whittle down the universe
/// to a managable size.
approx_filter: Option<PartitionedRibbonFilter<W, T, Approximate>>,
/// An exact membership query filter to confirm membership in R for items that
/// pass through the approximate filter.
exact_filter: Option<PartitionedRibbonFilter<W, T, Exact>>,
}
impl<const W: usize, T: Filterable<W>> Default for ClubcardBuilder<W, T> {
fn default() -> Self {
ClubcardBuilder {
approx_filter: None,
exact_filter: None,
}
}
}
impl<const W: usize, T: Filterable<W>> ClubcardBuilder<W, T> {
pub fn new() -> Self {
ClubcardBuilder::default()
}
pub fn new_approx_builder(&self, block: &[u8]) -> RibbonBuilder<'static, W, T> {
assert!(self.approx_filter.is_none());
RibbonBuilder::new(block, None)
}
pub fn new_exact_builder<'a>(&'a self, block: &[u8]) -> RibbonBuilder<'a, W, T> {
RibbonBuilder::new(block, self.approx_filter.as_ref())
}
pub fn collect_approx_ribbons(&mut self, ribbons: Vec<ApproximateRibbon<W, T>>) {
self.approx_filter = Some(PartitionedRibbonFilter::from(ribbons));
}
pub fn collect_exact_ribbons(&mut self, ribbons: Vec<Ribbon<W, T, Exact>>) {
self.exact_filter = Some(PartitionedRibbonFilter::from(ribbons));
}
pub fn build<U: Queryable<W>>(
self,
universe: U::UniverseMetadata,
partition: U::PartitionMetadata,
) -> Clubcard<W, U::UniverseMetadata, U::PartitionMetadata> {
let mut index: ClubcardIndex = BTreeMap::new();
assert!(self.approx_filter.is_some());
let approx_filter = self.approx_filter.unwrap();
for (block, entry) in approx_filter.index {
let meta = ClubcardIndexEntry {
approx_filter_offset: entry.offset,
approx_filter_m: entry.m,
approx_filter_rank: entry.rank,
exact_filter_offset: 0,
exact_filter_m: 0,
inverted: entry.inverted,
exceptions: entry.exceptions,
};
index.insert(block, meta);
}
assert!(self.exact_filter.is_some());
let mut exact_filter = self.exact_filter.unwrap();
for (block, entry) in exact_filter.index {
assert!(entry.rank == 1);
let meta = index.get_mut(&block).unwrap();
meta.exact_filter_offset = entry.offset;
meta.exact_filter_m = entry.m;
assert!(meta.inverted == entry.inverted);
meta.exceptions.extend(entry.exceptions);
}
assert!(exact_filter.solution.len() == 1);
let exact_filter = exact_filter.solution.pop().unwrap();
Clubcard {
universe,
partition,
index,
approx_filter: approx_filter.solution,
exact_filter,
}
}
}
#[cfg(test)]
mod tests {
use crate::builder::*;
use crate::*;
use rand::distributions::{Distribution, Uniform};
use rand::Rng;
// Construct the equation a(x) = x_i
fn std_eq<const W: usize>(i: usize) -> Equation<W> {
let mut a = [0u64; W];
a[0] = 1;
Equation::homogeneous(i, a)
}
// Construct an random aligned equation using the given distribution for s.
fn rand<const W: usize>(s_dist: &impl Distribution<usize>) -> Equation<W> {
let mut rng = rand::thread_rng();
let s = s_dist.sample(&mut rng);
let mut a = [0u64; W];
for a_i in a.iter_mut() {
*a_i = rng.gen();
}
a[0] |= 1;
Equation::inhomogeneous(s, a, rng.gen::<u8>() & 1)
}
impl<const W: usize> AsQuery<W> for Equation<W> {
fn as_query(&self, _m: usize) -> Equation<W> {
self.clone()
}
fn block(&self) -> &[u8] {
&[]
}
fn discriminant(&self) -> &[u8] {
unsafe { std::mem::transmute(&self.a[..]) }
}
}
impl<const W: usize> Filterable<W> for Equation<W> {
fn included(&self) -> bool {
self.b == 0
}
}
impl<const W: usize> Queryable<W> for Equation<W> {
type UniverseMetadata = ();
type PartitionMetadata = ();
fn in_universe(&self, _meta: &Self::UniverseMetadata) -> bool {
true
}
}
#[test]
fn test_solve_identity() {
let n = 1024;
let mut builder = RibbonBuilder::new(&[], None);
for i in 0usize..n {
let eq: Equation<1> = std_eq(i);
builder.insert(eq);
}
let ribbon = ExactRibbon::from(builder);
let filter = PartitionedRibbonFilter::from(vec![ribbon]);
for i in 0usize..n {
let eq: Equation<1> = std_eq(i);
assert!(eq.eval(&filter.solution[0]) == 0);
}
}
#[test]
fn test_solve_empty() {
let builder = RibbonBuilder::<4, Equation<4>>::new(&[0], None);
let ribbon = ApproximateRibbon::from(builder);
let filter = PartitionedRibbonFilter::from(vec![ribbon]);
assert!(!filter.contains(&std_eq(0)));
}
#[test]
fn test_solve_random() {
let n = 1024;
const W: usize = 2;
let mut r = Ribbon::<W, Equation<W>, Exact>::new(&[0], n, false);
let mut s_dist = Uniform::new(0, r.m);
let mut eqs = Vec::with_capacity(n);
for _ in 0..n {
let eq = rand(&mut s_dist);
eqs.push(eq.clone());
r.insert(eq);
}
let x = r.solve(&[]);
for eq in &eqs {
assert!(eq.eval(&x) == eq.b);
}
}
#[test]
fn test_total_approx_filter() {
// test that approximate filters that encode R=U are encoded
// as a zero-length solution vector with m=0 and inverted=true
// in the metadata.
let n = 1024;
let mut approx_builder = RibbonBuilder::new(&[], None);
approx_builder.set_universe_size(n);
for i in 0usize..n {
let eq: Equation<1> = std_eq(i);
approx_builder.insert(eq);
}
let approx_ribbon = ApproximateRibbon::from(approx_builder);
let approx_filter = PartitionedRibbonFilter::from(vec![approx_ribbon]);
let approx_index_entry = approx_filter
.index
.get(&vec![])
.expect("should have metadata");
assert!(approx_index_entry.m == 0);
assert!(approx_index_entry.rank == 0);
assert!(approx_index_entry.exceptions.is_empty());
assert!(approx_index_entry.inverted);
for i in 0usize..n {
let eq = std_eq(i);
assert!(approx_filter.contains(&eq));
}
assert!(approx_filter.solution.len() == 0);
let mut exact_builder = RibbonBuilder::new(&[], Some(&approx_filter));
for i in 0usize..n {
let mut eq = std_eq(i);
eq.b = 0;
exact_builder.insert(eq);
}
let exact_ribbon = ExactRibbon::from(exact_builder);
let exact_filter = PartitionedRibbonFilter::from(vec![exact_ribbon]);
let exact_index_entry = exact_filter
.index
.get(&vec![])
.expect("should have metadata");
assert!(exact_index_entry.m == 0);
assert!(exact_index_entry.rank == 1);
assert!(exact_index_entry.exceptions.is_empty());
assert!(exact_index_entry.inverted);
for i in 0usize..n {
let eq = std_eq(i);
assert!(exact_filter.contains(&eq));
}
assert!(exact_filter.solution.len() == 1);
assert!(exact_filter.solution[0].len() == 0);
}
#[test]
fn test_rank_0_approx_filter() {
let n = 1024;
let mut builder = RibbonBuilder::new(&[], None);
builder.set_universe_size(n);
for i in 0usize..768 {
let eq: Equation<1> = std_eq(i);
builder.insert(eq);
}
let ribbon = ApproximateRibbon::from(builder);
let filter = PartitionedRibbonFilter::from(vec![ribbon]);
let entry = filter.index.get(&vec![]).expect("should have metadata");
assert!(entry.rank == 0);
assert!(!entry.inverted);
assert!(filter.solution.len() == 0);
for i in 0usize..n {
let eq = std_eq(i);
assert!(filter.contains(&eq));
}
}
}