Merge pull request #7 from columbia/histogram_and_requests

Histogram and requests

Author: tholop

.github/workflows/rust.yml

@@ -21,4 +21,6 @@ jobs:
     - name: Run tests
       run: cargo test --verbose
     - name: Run integration tests
-      run: cargo test --package pdslib --test demo -- --nocapture 
+      run: |
+        cargo test --package pdslib --test ara_demo -- --nocapture
+        cargo test --package pdslib --test simple_events_demo -- --nocapture

justfile

@@ -7,7 +7,8 @@ test:
     cargo test
 
 demo:
-    cargo test --package pdslib --test demo -- --nocapture 
+    cargo test --package pdslib --test simple_events_demo -- --nocapture 
+    cargo test --package pdslib --test ara_demo -- --nocapture 
 
 format:
     cargo +nightly fmt

src/budget/hashmap_filter_storage.rs

@@ -1,4 +1,6 @@
-use crate::budget::traits::{Filter, FilterStorage, FilterStorageError};
+use crate::budget::traits::{
+    Budget, Filter, FilterStorage, FilterStorageError,
+};
 use std::collections::HashMap;
 use std::marker::PhantomData;
 
@@ -19,18 +21,19 @@ impl<K, F, Budget> HashMapFilterStorage<K, F, Budget> {
     }
 }
 
-impl<K, F, Budget> FilterStorage for HashMapFilterStorage<K, F, Budget>
+impl<K, F, B> FilterStorage for HashMapFilterStorage<K, F, B>
 where
-    F: Filter<Budget>,
+    B: Budget,
+    F: Filter<B>,
     K: Eq + std::hash::Hash,
 {
     type FilterId = K;
-    type Budget = Budget;
+    type Budget = B;
 
     fn new_filter(
         &mut self,
         filter_id: K,
-        capacity: Budget,
+        capacity: B,
     ) -> Result<(), FilterStorageError> {
         let filter = F::new(capacity);
         self.filters.insert(filter_id, filter);
@@ -45,7 +48,7 @@ where
     fn try_consume(
         &mut self,
         filter_id: &K,
-        budget: &Budget,
+        budget: &B,
     ) -> Result<(), FilterStorageError> {
         let filter = self
             .filters
@@ -54,6 +57,17 @@ where
         filter.try_consume(budget)?;
         Ok(())
     }
+
+    fn get_remaining_budget(
+        &self,
+        filter_id: &Self::FilterId,
+    ) -> Result<Self::Budget, FilterStorageError> {
+        let filter = self
+            .filters
+            .get(filter_id)
+            .ok_or(FilterStorageError::FilterDoesNotExist)?;
+        Ok(filter.get_remaining_budget())
+    }
 }
 
 #[cfg(test)]

src/budget/pure_dp_filter.rs

@@ -1,10 +1,12 @@
-pub use crate::budget::traits::{Filter, FilterError};
+pub use crate::budget::traits::{Budget, Filter, FilterError};
 
 #[derive(Debug, Clone)]
 pub struct PureDPBudget {
     pub epsilon: f64,
 }
 
+impl Budget for PureDPBudget {}
+
 // TODO: Check whether we can reuse the OpenDP accountant if we want to use RDP/zCDP, without having to execute a measurement on real data. Check out the `compose` function here: https://docs.rs/opendp/latest/opendp/measures/struct.ZeroConcentratedDivergence.html, check if they offer filters directly.
 
 #[derive(Debug)]
@@ -31,6 +33,10 @@ impl Filter<PureDPBudget> for PureDPBudgetFilter {
             Err(FilterError::OutOfBudget)
         }
     }
+
+    fn get_remaining_budget(&self) -> PureDPBudget {
+        self.remaining_budget.clone()
+    }
 }
 
 #[cfg(test)]

src/budget/traits.rs

@@ -1,7 +1,10 @@
-// TODO: maybe Budget trait, and Filter<T: Budget> if we need?
-
 use thiserror::Error;
 
+/// Trait for privacy budgets
+pub trait Budget: Clone {
+    // For now just a marker trait requiring Clone
+}
+
 /// Error returned when trying to consume from a filter.
 #[derive(Error, Debug)]
 pub enum FilterError {
@@ -10,12 +13,18 @@ pub enum FilterError {
 }
 
 /// Trait for a privacy filter.
-pub trait Filter<T> {
+pub trait Filter<T: Budget> {
     /// Initializes a new filter with a given capacity.
     fn new(capacity: T) -> Self;
 
-    /// Tries to consume a given budget from the filter. In the formalism from https://arxiv.org/abs/1605.08294, Ok(()) corresponds to CONTINUE, and Err(FilterError::OutOfBudget) corresponds to HALT..
+    /// Tries to consume a given budget from the filter.
+    /// In the formalism from https://arxiv.org/abs/1605.08294, Ok(()) corresponds to CONTINUE, and Err(FilterError::OutOfBudget) corresponds to HALT.
     fn try_consume(&mut self, budget: &T) -> Result<(), FilterError>;
+
+    /// Gets the remaining budget for this filter.
+    /// WARNING: this method is for local visualization only.
+    /// Its output should not be shared outside the device.
+    fn get_remaining_budget(&self) -> T;
 }
 
 /// Error returned when trying to interact with a filter storage.
@@ -32,7 +41,7 @@ pub enum FilterStorageError {
 /// Trait for an interface or object that maintains a collection of filters.
 pub trait FilterStorage {
     type FilterId;
-    type Budget;
+    type Budget: Budget;
 
     /// Initializes a new filter with an associated filter ID and capacity.
     fn new_filter(
@@ -50,4 +59,10 @@ pub trait FilterStorage {
         filter_id: &Self::FilterId,
         budget: &Self::Budget,
     ) -> Result<(), FilterStorageError>;
+
+    /// Gets the remaining budget for a filter.
+    fn get_remaining_budget(
+        &self,
+        filter_id: &Self::FilterId,
+    ) -> Result<Self::Budget, FilterStorageError>;
 }

src/events/ara_event.rs

@@ -0,0 +1,23 @@
+use std::collections::HashMap;
+
+use crate::events::traits::Event;
+
+/// Source event for ARA-style callers such as Chromium.
+/// Mimics the fields from https://source.chromium.org/chromium/chromium/src/+/main:content/browser/attribution_reporting/attribution_reporting.proto.
+/// TODO: add other fields as needed by callers.
+/// TODO: replace String by usize if that simplifies FFI?
+#[derive(Debug, Clone)]
+pub struct AraEvent {
+    pub id: usize,
+    pub epoch_number: usize,
+    pub aggregatable_sources: HashMap<String, usize>,
+    // TODO: add filters here
+}
+
+impl Event for AraEvent {
+    type EpochId = usize;
+
+    fn get_epoch_id(&self) -> Self::EpochId {
+        self.epoch_number
+    }
+}

src/events/hashmap_event_storage.rs

@@ -0,0 +1,67 @@
+use crate::events::traits::EventStorage;
+use crate::events::traits::RelevantEventSelector;
+use crate::events::traits::{EpochEvents, Event};
+
+use std::collections::HashMap;
+use std::marker::PhantomData;
+
+pub type VecEpochEvents<E> = Vec<E>;
+
+impl<E: Event> EpochEvents for VecEpochEvents<E> {
+    fn is_empty(&self) -> bool {
+        self.is_empty()
+    }
+}
+
+/// A simple in-memory event storage. Stores a mapping of epoch id to epoch
+/// events, where each epoch events is just a vec of events.
+/// Clones events when asked to retrieve events for an epoch.
+#[derive(Debug)]
+pub struct HashMapEventStorage<E: Event, RES: RelevantEventSelector<Event = E>>
+{
+    epochs: HashMap<E::EpochId, VecEpochEvents<E>>,
+    _phantom: PhantomData<RES>,
+}
+
+impl<E: Event, RES: RelevantEventSelector<Event = E>>
+    HashMapEventStorage<E, RES>
+{
+    pub fn new() -> Self {
+        Self {
+            epochs: HashMap::new(),
+            _phantom: PhantomData,
+        }
+    }
+}
+
+impl<E, RES> EventStorage for HashMapEventStorage<E, RES>
+where
+    E: Event + Clone,
+    RES: RelevantEventSelector<Event = E>,
+{
+    type Event = E;
+    type EpochEvents = VecEpochEvents<E>;
+    type RelevantEventSelector = RES;
+
+    fn add_event(&mut self, event: E) -> Result<(), ()> {
+        let epoch_id = event.get_epoch_id();
+        let epoch = self.epochs.entry(epoch_id).or_default();
+        epoch.push(event);
+        Ok(())
+    }
+
+    fn get_epoch_events(
+        &self,
+        epoch_id: &E::EpochId,
+        selector: &RES,
+    ) -> Option<VecEpochEvents<E>> {
+        // Return relevant events for a given epoch_id
+        self.epochs.get(&epoch_id).map(|events| {
+            events
+                .iter()
+                .filter(|event| selector.is_relevant_event(event))
+                .cloned()
+                .collect()
+        })
+    }
+}

src/events/mod.rs

@@ -1,2 +1,4 @@
-pub mod simple_events;
+pub mod ara_event;
+pub mod hashmap_event_storage;
+pub mod simple_event;
 pub mod traits;

src/events/simple_event.rs

@@ -0,0 +1,34 @@
+use crate::events::traits::Event;
+
+/// TODO: add enough things to run basic queries and filter by attributes.
+/// use https://github.com/patcg/meetings/blob/main/2024/09/27-tpac/Privacy-Preserving%20Attribution%20Proposed%20Roadmap.pdf
+#[derive(Debug, Clone)]
+pub struct SimpleEvent {
+    pub id: usize,
+    pub epoch_number: usize,
+    pub event_key: usize,
+    // TODO: consider adding timestamp
+}
+
+impl Event for SimpleEvent {
+    type EpochId = usize;
+
+    fn get_epoch_id(&self) -> Self::EpochId {
+        self.epoch_number
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_simple_event() {
+        let event = SimpleEvent {
+            id: 1,
+            epoch_number: 1,
+            event_key: 3,
+        };
+        assert_eq!(event.id, 1);
+    }
+}