oram: add basic tracing of runtimes

oram/Cargo.toml

@@ -18,6 +18,8 @@ num-bigint = "0.4.3"
 num-traits = "0.2.15"
 rand = "0.8.5"
 rand_chacha = "0.3.1"
+strum = { version = "0.24.1", features = ["derive"] }
+strum_macros = "0.24"
 
 [dev-dependencies]
 criterion = "0.4.0"

oram/src/oram.rs

@@ -2,7 +2,9 @@ use crate::common::{Error, InstructionShare};
 use crate::doprf::{JointDOPrf, LegendrePrf, LegendrePrfKey};
 use crate::p_ot::{POTIndexParty, POTKeyParty, POTReceiverParty};
 use crate::select::{Select, SelectProtocol};
-use crate::stash::{Stash, StashProtocol};
+use crate::stash::{
+    ProtocolStep as StashProtocolStep, Runtimes as StashRuntimes, Stash, StashProtocol,
+};
 use communicator::{AbstractCommunicator, Fut, Serializable};
 use dpf::{mpdpf::MultiPointDpf, spdpf::SinglePointDpf};
 use ff::PrimeField;
@@ -10,6 +12,8 @@ use itertools::{izip, Itertools};
 use rand::thread_rng;
 use std::iter::repeat;
 use std::marker::PhantomData;
+use std::time::{Duration, Instant};
+use strum::IntoEnumIterator;
 use utils::field::{FromPrf, LegendreSymbol};
 use utils::permutation::FisherYatesPermutation;
 
@@ -44,6 +48,113 @@ fn compute_oram_prf_output_bitsize(memory_size: usize) -> usize {
     (usize::BITS - memory_size.leading_zeros()) as usize + 40
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq, strum_macros::EnumIter, strum_macros::Display)]
+pub enum ProtocolStep {
+    AccessStashRead = 0,
+    AccessAddressSelection,
+    AccessDatabaseRead,
+    AccessStashWrite,
+    AccessValueSelection,
+    AccessRefresh,
+    DbReadAddressTag,
+    DbReadGarbledIndex,
+    DbReadPotAccess,
+    DbWriteMpDpfKeyExchange,
+    DbWriteMpDpfEvaluations,
+    DbWriteUpdateMemory,
+    RefreshResetFuncs,
+    RefreshInitStash,
+    RefreshInitDOPrf,
+    RefreshInitPOt,
+    RefreshGarbleMemory,
+}
+
+#[derive(Debug, Default, Clone, Copy)]
+pub struct Runtimes {
+    durations: [Duration; 17],
+    stash_runtimes: StashRuntimes,
+}
+
+impl Runtimes {
+    #[inline(always)]
+    pub fn record(&mut self, id: ProtocolStep, duration: Duration) {
+        self.durations[id as usize] += duration;
+    }
+
+    pub fn get_stash_runtimes(&self) -> StashRuntimes {
+        self.stash_runtimes
+    }
+
+    pub fn set_stash_runtimes(&mut self, stash_runtimes: StashRuntimes) {
+        self.stash_runtimes = stash_runtimes;
+    }
+
+    pub fn get(&self, id: ProtocolStep) -> Duration {
+        self.durations[id as usize]
+    }
+
+    pub fn print(&self, party_id: usize, num_accesses: usize) {
+        println!("=============== Party {} ===============", party_id);
+        println!("- times per access over {num_accesses} accesses in total");
+        for step in ProtocolStep::iter().filter(|x| x.to_string().starts_with("Access")) {
+            println!(
+                "{:28}    {:.5} s",
+                step,
+                self.get(step).as_secs_f64() / num_accesses as f64
+            );
+            match step {
+                ProtocolStep::AccessDatabaseRead => {
+                    for step in ProtocolStep::iter().filter(|x| x.to_string().starts_with("DbRead"))
+                    {
+                        println!(
+                            "    {:24}    {:.5} s",
+                            step,
+                            self.get(step).as_secs_f64() / num_accesses as f64
+                        );
+                    }
+                }
+                ProtocolStep::AccessRefresh => {
+                    for step in ProtocolStep::iter().filter(|x| {
+                        x.to_string().starts_with("DbWrite") || x.to_string().starts_with("Refresh")
+                    }) {
+                        println!(
+                            "    {:24}    {:.5} s",
+                            step,
+                            self.get(step).as_secs_f64() / num_accesses as f64
+                        );
+                    }
+                }
+                ProtocolStep::AccessStashRead => {
+                    for step in
+                        StashProtocolStep::iter().filter(|x| x.to_string().starts_with("Read"))
+                    {
+                        println!(
+                            "    {:24}    {:.5} s",
+                            step,
+                            self.stash_runtimes.get(step).as_secs_f64() / num_accesses as f64
+                        );
+                    }
+                }
+                ProtocolStep::AccessStashWrite => {
+                    for step in
+                        StashProtocolStep::iter().filter(|x| x.to_string().starts_with("Write"))
+                    {
+                        println!(
+                            "    {:24}    {:.5} s",
+                            step,
+                            self.stash_runtimes.get(step).as_secs_f64() / num_accesses as f64
+                        );
+                    }
+                }
+                _ => {}
+            }
+            // if step.to_string().starts_with("") {
+            // }
+        }
+        println!("========================================");
+    }
+}
+
 pub struct DistributedOramProtocol<F, MPDPF, SPDPF>
 where
     F: FromPrf + LegendreSymbol + Serializable,
@@ -148,30 +259,61 @@ where
         &mut self,
         comm: &mut C,
         address_share: F,
-    ) -> Result<F, Error> {
+        runtimes: Option<Runtimes>,
+    ) -> Result<(F, Option<Runtimes>), Error> {
         let mut value_share = F::ZERO;
 
+        let t_start = Instant::now();
+
         // 1. Compute address tag
         let address_tag: u128 = self.joint_doprf.eval_to_uint(comm, &[address_share])?[0];
 
         // 2. Update tags read list
         self.address_tags_read.push(address_tag);
 
+        let t_after_address_tag = Instant::now();
+
         // 3. Compute index in garbled memory and retrieve share
         let garbled_index = self.pos_mine(address_tag);
         value_share += self.garbled_memory_share[garbled_index].1;
 
+        let t_after_index_computation = Instant::now();
+
         // 4. Run p-OT.Access
         self.pot_index_party.run_access(comm, garbled_index)?;
         value_share -= self.pot_receiver_party.run_access(comm)?;
 
-        Ok(value_share)
+        let t_after_pot_access = Instant::now();
+
+        let runtimes = runtimes.map(|mut r| {
+            r.record(
+                ProtocolStep::DbReadAddressTag,
+                t_after_address_tag - t_start,
+            );
+            r.record(
+                ProtocolStep::DbReadAddressTag,
+                t_after_index_computation - t_after_address_tag,
+            );
+            r.record(
+                ProtocolStep::DbReadAddressTag,
+                t_after_pot_access - t_after_index_computation,
+            );
+            r
+        });
+
+        Ok((value_share, runtimes))
     }
 
     fn update_database_from_stash<C: AbstractCommunicator>(
         &mut self,
         comm: &mut C,
-    ) -> Result<(), Error> {
+        runtimes: Option<Runtimes>,
+    ) -> Result<Option<Runtimes>, Error> {
+        let t_start = Instant::now();
+
+        let fut_dpf_key_from_prev = comm.receive_previous()?;
+        let fut_dpf_key_from_next = comm.receive_next()?;
+
         let mpdpf = MPDPF::new(self.memory_size, self.stash_size);
         let mut points = Vec::with_capacity(self.stash_size);
         let mut values = Vec::with_capacity(self.stash_size);
@@ -198,17 +340,19 @@ where
             (points, new_values)
         };
 
-        let fut_dpf_key_from_prev = comm.receive_previous()?;
-        let fut_dpf_key_from_next = comm.receive_next()?;
         let (dpf_key_prev, dpf_key_next) = mpdpf.generate_keys(&points, &values);
         comm.send_previous(dpf_key_prev)?;
         comm.send_next(dpf_key_next)?;
         let dpf_key_from_prev = fut_dpf_key_from_prev.get()?;
         let dpf_key_from_next = fut_dpf_key_from_next.get()?;
 
+        let t_after_mpdpf_key_exchange = Instant::now();
+
         let new_memory_share_from_prev = mpdpf.evaluate_domain(&dpf_key_from_prev);
         let new_memory_share_from_next = mpdpf.evaluate_domain(&dpf_key_from_next);
 
+        let t_after_mpdpf_evaluations = Instant::now();
+
         {
             let mut memory_share = Vec::new();
             std::mem::swap(&mut self.memory_share, &mut memory_share);
@@ -220,10 +364,34 @@ where
             std::mem::swap(&mut self.memory_share, &mut memory_share);
         }
 
-        Ok(())
+        let t_after_memory_update = Instant::now();
+
+        let runtimes = runtimes.map(|mut r| {
+            r.record(
+                ProtocolStep::DbWriteMpDpfKeyExchange,
+                t_after_mpdpf_key_exchange - t_start,
+            );
+            r.record(
+                ProtocolStep::DbWriteMpDpfEvaluations,
+                t_after_mpdpf_evaluations - t_after_mpdpf_key_exchange,
+            );
+            r.record(
+                ProtocolStep::DbWriteUpdateMemory,
+                t_after_memory_update - t_after_mpdpf_evaluations,
+            );
+            r
+        });
+
+        Ok(runtimes)
     }
 
-    fn refresh<C: AbstractCommunicator>(&mut self, comm: &mut C) -> Result<(), Error> {
+    fn refresh<C: AbstractCommunicator>(
+        &mut self,
+        comm: &mut C,
+        runtimes: Option<Runtimes>,
+    ) -> Result<Option<Runtimes>, Error> {
+        let t_start = Instant::now();
+
         // 0. Reset the functionalities
         self.stash.reset();
         self.joint_doprf.reset();
@@ -231,9 +399,13 @@ where
         self.pot_index_party.reset();
         self.pot_receiver_party.reset();
 
+        let t_after_reset = Instant::now();
+
         // 1. Initialize the stash
         self.stash.init(comm)?;
 
+        let t_after_init_stash = Instant::now();
+
         // 2. Run r-DB init protocol
         // a) Initialize DOPRF
         {
@@ -247,6 +419,8 @@ where
             self.joint_doprf.preprocess(comm, self.stash_size)?;
         }
 
+        let t_after_init_doprf = Instant::now();
+
         // b) Initialize p-OT
         {
             match self.party_id {
@@ -269,7 +443,10 @@ where
             };
         }
 
+        let t_after_init_pot = Instant::now();
+
         // c) Compute index tags and garble the memory share for the next party
+        let fut_garbled_memory_share = comm.receive_previous()?;
         self.memory_index_tags_prev = Vec::with_capacity(self.memory_size);
         self.memory_index_tags_prev
             .extend((0..self.memory_size).map(|j| {
@@ -290,7 +467,6 @@ where
                 .all(|w| w[0] < w[1]),
             "index tags not sorted or colliding"
         );
-        let fut_garbled_memory_share = comm.receive_previous()?;
         let mut garbled_memory_share_next: Vec<_> = self
             .memory_share
             .iter()
@@ -331,7 +507,119 @@ where
         // the garbled_memory_share now defines the pos_mine map:
         // - pos_i(tag) -> index of tag in garbled_memory_share
 
-        Ok(())
+        let t_after_garble_memory = Instant::now();
+
+        let runtimes = runtimes.map(|mut r| {
+            r.record(ProtocolStep::RefreshResetFuncs, t_after_reset - t_start);
+            r.record(
+                ProtocolStep::RefreshInitStash,
+                t_after_init_stash - t_after_reset,
+            );
+            r.record(
+                ProtocolStep::RefreshInitDOPrf,
+                t_after_init_doprf - t_after_init_stash,
+            );
+            r.record(
+                ProtocolStep::RefreshInitPOt,
+                t_after_init_pot - t_after_init_doprf,
+            );
+            r.record(
+                ProtocolStep::RefreshGarbleMemory,
+                t_after_garble_memory - t_after_init_pot,
+            );
+            r
+        });
+
+        Ok(runtimes)
+    }
+
+    pub fn access_with_runtimes<C: AbstractCommunicator>(
+        &mut self,
+        comm: &mut C,
+        instruction: InstructionShare<F>,
+        runtimes: Option<Runtimes>,
+    ) -> Result<(F, Option<Runtimes>), Error> {
+        assert!(self.is_initialized);
+
+        // 1. Read from the stash
+        let t_start = Instant::now();
+        let (stash_state, stash_runtimes) = self.stash.read_with_runtimes(
+            comm,
+            instruction,
+            runtimes.map(|r| r.get_stash_runtimes()),
+        )?;
+        let t_after_stash_read = Instant::now();
+
+        // 2. If the value was found in a stash, we read from the dummy address
+        let dummy_address_share = match self.party_id {
+            PARTY_1 => F::from_u128((1 << self.log_db_size) + self.get_access_counter() as u128),
+            _ => F::ZERO,
+        };
+        let db_address_share = SelectProtocol::select(
+            comm,
+            stash_state.flag,
+            dummy_address_share,
+            instruction.address,
+        )?;
+        let t_after_address_selection = Instant::now();
+
+        // 3. Read a (dummy or real) value from the database
+        let (db_value_share, runtimes) =
+            self.read_from_database(comm, db_address_share, runtimes)?;
+        let t_after_db_read = Instant::now();
+
+        // 4. Write the read value into the stash
+        let stash_runtime = self.stash.write_with_runtimes(
+            comm,
+            instruction,
+            stash_state,
+            db_address_share,
+            db_value_share,
+            stash_runtimes,
+        )?;
+        let t_after_stash_write = Instant::now();
+
+        // 5. Select the right value to return
+        let read_value =
+            SelectProtocol::select(comm, stash_state.flag, stash_state.value, db_value_share)?;
+        let t_after_value_selection = Instant::now();
+
+        // 6. If the stash is full, write the value back into the database
+        let runtimes = if self.get_access_counter() == self.stash_size {
+            let runtimes = self.update_database_from_stash(comm, runtimes)?;
+            self.refresh(comm, runtimes)?
+        } else {
+            runtimes
+        };
+        let t_after_refresh = Instant::now();
+
+        let runtimes = runtimes.map(|mut r| {
+            r.set_stash_runtimes(stash_runtime.unwrap());
+            r.record(ProtocolStep::AccessStashRead, t_after_stash_read - t_start);
+            r.record(
+                ProtocolStep::AccessAddressSelection,
+                t_after_address_selection - t_after_stash_read,
+            );
+            r.record(
+                ProtocolStep::AccessDatabaseRead,
+                t_after_db_read - t_after_address_selection,
+            );
+            r.record(
+                ProtocolStep::AccessStashWrite,
+                t_after_stash_write - t_after_db_read,
+            );
+            r.record(
+                ProtocolStep::AccessValueSelection,
+                t_after_value_selection - t_after_stash_write,
+            );
+            r.record(
+                ProtocolStep::AccessRefresh,
+                t_after_refresh - t_after_value_selection,
+            );
+            r
+        });
+
+        Ok((read_value, runtimes))
     }
 }
 
@@ -363,7 +651,7 @@ where
             .extend(repeat(F::ZERO).take(self.stash_size));
 
         // 2. Run the refresh protocol to initialize everything.
-        self.refresh(comm)?;
+        self.refresh(comm, None)?;
 
         self.is_initialized = true;
         Ok(())
@@ -374,46 +662,8 @@ where
         comm: &mut C,
         instruction: InstructionShare<F>,
     ) -> Result<F, Error> {
-        assert!(self.is_initialized);
-
-        // 1. Read from the stash
-        let stash_state = self.stash.read(comm, instruction)?;
-
-        // 2. If the value was found in a stash, we read from the dummy address
-        let dummy_address_share = match self.party_id {
-            PARTY_1 => F::from_u128((1 << self.log_db_size) + self.get_access_counter() as u128),
-            _ => F::ZERO,
-        };
-        let db_address_share = SelectProtocol::select(
-            comm,
-            stash_state.flag,
-            dummy_address_share,
-            instruction.address,
-        )?;
-
-        // 3. Read a (dummy or real) value from the database
-        let db_value_share = self.read_from_database(comm, db_address_share)?;
-
-        // 4. Write the read value into the stash
-        self.stash.write(
-            comm,
-            instruction,
-            stash_state,
-            db_address_share,
-            db_value_share,
-        )?;
-
-        // 5. Select the right value to return
-        let read_value =
-            SelectProtocol::select(comm, stash_state.flag, stash_state.value, db_value_share)?;
-
-        // 6. If the stash is full, write the value back into the database
-        if self.get_access_counter() == self.stash_size {
-            self.update_database_from_stash(comm)?;
-            self.refresh(comm)?;
-        }
-
-        Ok(read_value)
+        self.access_with_runtimes(comm, instruction, None)
+            .map(|x| x.0)
     }
 
     fn get_db<C: AbstractCommunicator>(
@@ -424,7 +674,7 @@ where
         assert!(self.is_initialized);
 
         if self.get_access_counter() > 0 {
-            self.refresh(comm)?;
+            self.refresh(comm, None)?;
         }
 
         if rerandomize_shares {

oram/src/stash.rs

@@ -10,6 +10,7 @@ use dpf::spdpf::SinglePointDpf;
 use ff::PrimeField;
 use rand::thread_rng;
 use std::marker::PhantomData;
+use std::time::{Duration, Instant};
 use utils::field::LegendreSymbol;
 
 #[derive(Clone, Copy, Debug, Default)]
@@ -71,97 +72,41 @@ fn compute_stash_prf_output_bitsize(stash_size: usize) -> usize {
     (usize::BITS - stash_size.leading_zeros()) as usize + 40
 }
 
-fn stash_read_value<C, F, SPDPF>(
-    comm: &mut C,
-    access_counter: usize,
-    location_share: F,
-    stash_values_share_mine: &[F],
-) -> Result<F, Error>
-where
-    C: AbstractCommunicator,
-    F: PrimeField + Serializable,
-    SPDPF: SinglePointDpf<Value = F>,
-    SPDPF::Key: Serializable,
-{
-    // a) convert the stash into replicated secret sharing
-    let fut_prev = comm.receive_previous::<Vec<F>>()?;
-    comm.send_next(stash_values_share_mine.to_vec())?;
-    let stash_values_share_prev = fut_prev.get()?;
-
-    // b) mask and reconstruct the stash index <loc>
-    let index_bits = (access_counter as f64).log2().ceil() as u32;
-    assert!(index_bits <= 16);
-    let bit_mask = ((1 << index_bits) - 1) as u16;
-    let (masked_loc, r_prev, r_next) =
-        MaskIndexProtocol::mask_index(comm, index_bits, location_share)?;
-
-    // c) use DPFs to read the stash value
-    let fut_prev = comm.receive_previous::<SPDPF::Key>()?;
-    let fut_next = comm.receive_next::<SPDPF::Key>()?;
-    {
-        let (dpf_key_prev, dpf_key_next) =
-            SPDPF::generate_keys(1 << index_bits, masked_loc as u64, F::ONE);
-        comm.send_previous(dpf_key_prev)?;
-        comm.send_next(dpf_key_next)?;
-    }
-    let dpf_key_prev = fut_prev.get()?;
-    let dpf_key_next = fut_next.get()?;
-    let mut value_share = F::ZERO;
-    for j in 0..access_counter {
-        let index_prev = ((j as u16 + r_prev) & bit_mask) as u64;
-        let index_next = ((j as u16 + r_next) & bit_mask) as u64;
-        value_share += SPDPF::evaluate_at(&dpf_key_prev, index_prev) * stash_values_share_mine[j];
-        value_share += SPDPF::evaluate_at(&dpf_key_next, index_next) * stash_values_share_prev[j];
-    }
-    Ok(value_share)
+#[derive(Debug, Clone, Copy, PartialEq, Eq, strum_macros::EnumIter, strum_macros::Display)]
+pub enum ProtocolStep {
+    Init = 0,
+    ReadMaskedAddressTag,
+    ReadDpfKeyGen,
+    ReadLookupFlagLocation,
+    ReadComputeLocation,
+    ReadReshareFlag,
+    ReadConvertToReplicated,
+    ReadComputeMaskedIndex,
+    ReadDpfKeyDistribution,
+    ReadDpfEvaluations,
+    WriteAddressTag,
+    WriteStoreTriple,
+    WriteSelectPreviousValue,
+    WriteSelectValue,
+    WriteComputeMaskedIndex,
+    WriteDpfKeyDistribution,
+    WriteDpfEvaluations,
 }
 
-fn stash_write_value<C, F, SPDPF>(
-    comm: &mut C,
-    access_counter: usize,
-    location_share: F,
-    // old_value_share: F,
-    value_share: F,
-    stash_values_share_mine: &mut [F],
-) -> Result<(), Error>
-where
-    C: AbstractCommunicator,
-    F: PrimeField + Serializable,
-    SPDPF: SinglePointDpf<Value = F>,
-    SPDPF::Key: Serializable,
-{
-    // a) mask and reconstruct the stash index <loc>
-    let index_bits = {
-        let bits = usize::BITS - access_counter.leading_zeros();
-        if bits > 0 {
-            bits
-        } else {
-            1
-        }
-    };
-    assert!(index_bits <= 16);
-    let bit_mask = ((1 << index_bits) - 1) as u16;
-    let (masked_loc, r_prev, r_next) =
-        MaskIndexProtocol::mask_index(comm, index_bits, location_share)?;
-
-    // b) use DPFs to read the stash value
-    let fut_prev = comm.receive_previous::<SPDPF::Key>()?;
-    let fut_next = comm.receive_next::<SPDPF::Key>()?;
-    {
-        let (dpf_key_prev, dpf_key_next) =
-            SPDPF::generate_keys(1 << index_bits, masked_loc as u64, value_share);
-        comm.send_previous(dpf_key_prev)?;
-        comm.send_next(dpf_key_next)?;
+#[derive(Debug, Default, Clone, Copy)]
+pub struct Runtimes {
+    durations: [Duration; 17],
+}
+
+impl Runtimes {
+    #[inline(always)]
+    pub fn record(&mut self, id: ProtocolStep, duration: Duration) {
+        self.durations[id as usize] += duration;
     }
-    let dpf_key_prev = fut_prev.get()?;
-    let dpf_key_next = fut_next.get()?;
-    for j in 0..=access_counter {
-        let index_prev = ((j as u16).wrapping_add(r_prev) & bit_mask) as u64;
-        let index_next = ((j as u16).wrapping_add(r_next) & bit_mask) as u64;
-        stash_values_share_mine[j] += SPDPF::evaluate_at(&dpf_key_prev, index_prev);
-        stash_values_share_mine[j] += SPDPF::evaluate_at(&dpf_key_next, index_next);
+
+    pub fn get(&self, id: ProtocolStep) -> Duration {
+        self.durations[id as usize]
     }
-    Ok(())
 }
 
 pub struct StashProtocol<F, SPDPF>
@@ -190,6 +135,7 @@ impl<F, SPDPF> StashProtocol<F, SPDPF>
 where
     F: PrimeField + LegendreSymbol + Serializable,
     SPDPF: SinglePointDpf<Value = F>,
+    SPDPF::Key: Serializable,
 {
     pub fn new(party_id: usize, stash_size: usize) -> Self {
         assert!(party_id < 3);
@@ -218,33 +164,16 @@ where
             _phantom: PhantomData,
         }
     }
-}
 
-impl<F, SPDPF> Stash<F> for StashProtocol<F, SPDPF>
-where
-    F: PrimeField + LegendreSymbol + Serializable,
-    SPDPF: SinglePointDpf<Value = F>,
-    SPDPF::Key: Serializable,
-{
-    fn get_party_id(&self) -> usize {
-        self.party_id
-    }
-
-    fn get_stash_size(&self) -> usize {
-        self.stash_size
-    }
-
-    fn get_access_counter(&self) -> usize {
-        self.access_counter
-    }
-
-    fn reset(&mut self) {
-        *self = Self::new(self.party_id, self.stash_size);
-    }
-
-    fn init<C: AbstractCommunicator>(&mut self, comm: &mut C) -> Result<(), Error> {
+    fn init_with_runtimes<C: AbstractCommunicator>(
+        &mut self,
+        comm: &mut C,
+        runtimes: Option<Runtimes>,
+    ) -> Result<Option<Runtimes>, Error> {
         assert_eq!(self.state, State::New);
 
+        let t_start = Instant::now();
+
         let prf_output_bitsize = compute_stash_prf_output_bitsize(self.stash_size);
         let legendre_prf_key = LegendrePrf::<F>::key_gen(prf_output_bitsize);
 
@@ -255,6 +184,8 @@ where
                 let mut mdoprf_p1 = MaskedDOPrfParty1::from_legendre_prf_key(legendre_prf_key);
                 doprf_p1.init(comm)?;
                 mdoprf_p1.init(comm)?;
+                doprf_p1.preprocess(comm, self.stash_size)?;
+                mdoprf_p1.preprocess(comm, self.stash_size)?;
                 self.doprf_party_1 = Some(doprf_p1);
                 self.masked_doprf_party_1 = Some(mdoprf_p1);
             }
@@ -263,6 +194,8 @@ where
                 let mut mdoprf_p2 = MaskedDOPrfParty2::new(prf_output_bitsize);
                 doprf_p2.init(comm)?;
                 mdoprf_p2.init(comm)?;
+                doprf_p2.preprocess(comm, self.stash_size)?;
+                mdoprf_p2.preprocess(comm, self.stash_size)?;
                 self.doprf_party_2 = Some(doprf_p2);
                 self.masked_doprf_party_2 = Some(mdoprf_p2);
             }
@@ -271,45 +204,65 @@ where
                 let mut mdoprf_p3 = MaskedDOPrfParty3::new(prf_output_bitsize);
                 doprf_p3.init(comm)?;
                 mdoprf_p3.init(comm)?;
+                doprf_p3.preprocess(comm, self.stash_size)?;
+                mdoprf_p3.preprocess(comm, self.stash_size)?;
                 self.doprf_party_3 = Some(doprf_p3);
                 self.masked_doprf_party_3 = Some(mdoprf_p3);
             }
             _ => panic!("invalid party id"),
         }
 
+        let t_end = Instant::now();
+        let runtimes = runtimes.map(|mut r| {
+            r.record(ProtocolStep::Init, t_end - t_start);
+            r
+        });
+
         // panic!("not implemented");
         self.state = State::AwaitingRead;
-        Ok(())
+        Ok(runtimes)
     }
 
-    fn read<C: AbstractCommunicator>(
+    pub fn read_with_runtimes<C: AbstractCommunicator>(
         &mut self,
         comm: &mut C,
         instruction: InstructionShare<F>,
-    ) -> Result<StashStateShare<F>, Error> {
+        runtimes: Option<Runtimes>,
+    ) -> Result<(StashStateShare<F>, Option<Runtimes>), Error> {
         assert_eq!(self.state, State::AwaitingRead);
         assert!(self.access_counter < self.stash_size);
 
         // 0. If the stash is empty, we are done
         if self.access_counter == 0 {
             self.state = State::AwaitingWrite;
-            return Ok(StashStateShare {
-                flag: F::ZERO,
-                location: F::ZERO,
-                value: F::ZERO,
-            });
+            return Ok((
+                StashStateShare {
+                    flag: F::ZERO,
+                    location: F::ZERO,
+                    value: F::ZERO,
+                },
+                runtimes,
+            ));
         }
 
-        let (flag_share, location_share) = match self.party_id {
+        let t_start = Instant::now();
+
+        let (
+            flag_share,
+            location_share,
+            t_after_masked_address_tag,
+            t_after_dpf_keygen,
+            t_after_compute_flag_loc,
+        ) = match self.party_id {
             PARTY_1 => {
                 // 1. Compute tag y := PRF(k, <I.adr>) such that P1 obtains y + r and P2, P3 obtain the mask r.
                 let masked_address_tag: u64 = {
                     let mdoprf_p1 = self.masked_doprf_party_1.as_mut().unwrap();
-                    // for now do preprocessing on the fly
-                    mdoprf_p1.preprocess(comm, 1)?;
                     mdoprf_p1.eval_to_uint(comm, 1, &[instruction.address])?[0]
                 };
 
+                let t_after_masked_address_tag = Instant::now();
+
                 // 2. Create and send DPF keys for the function f(x) = if x = y { 1 } else { 0 }
                 {
                     let domain_size = 1 << compute_stash_prf_output_bitsize(self.stash_size);
@@ -319,34 +272,42 @@ where
                     comm.send(PARTY_3, dpf_key_3)?;
                 }
 
+                let t_after_dpf_keygen = Instant::now();
+
                 // 3. The other parties compute shares of <flag>, <loc>, i.e., if the address is present in
                 //    the stash and if so, where it is. We just take 0s as our shares.
-                (F::ZERO, F::ZERO)
+                (
+                    F::ZERO,
+                    F::ZERO,
+                    t_after_masked_address_tag,
+                    t_after_dpf_keygen,
+                    t_after_dpf_keygen,
+                )
             }
             PARTY_2 | PARTY_3 => {
                 // 1. Compute tag y := PRF(k, <I.adr>) such that P1 obtains y + r and P2, P3 obtain the mask r.
                 let address_tag_mask: u64 = match self.party_id {
                     PARTY_2 => {
                         let mdoprf_p2 = self.masked_doprf_party_2.as_mut().unwrap();
-                        // for now do preprocessing on the fly
-                        mdoprf_p2.preprocess(comm, 1)?;
                         mdoprf_p2.eval_to_uint(comm, 1, &[instruction.address])?[0]
                     }
                     PARTY_3 => {
                         let mdoprf_p3 = self.masked_doprf_party_3.as_mut().unwrap();
-                        // for now do preprocessing on the fly
-                        mdoprf_p3.preprocess(comm, 1)?;
                         mdoprf_p3.eval_to_uint(comm, 1, &[instruction.address])?[0]
                     }
                     _ => panic!("invalid party id"),
                 };
 
+                let t_after_masked_address_tag = Instant::now();
+
                 // 2. Receive DPF key for the function f(x) = if x = y { 1 } else { 0 }
                 let dpf_key_i: SPDPF::Key = {
                     let fut = comm.receive(PARTY_1)?;
                     fut.get()?
                 };
 
+                let t_after_dpf_keygen = Instant::now();
+
                 // 3. Compute shares of <flag>, <loc>, i.e., if the address is present in the stash and if
                 //    so, where it is
                 {
@@ -362,7 +323,14 @@ where
                         location_share += j_as_field_element * dpf_value_j;
                         j_as_field_element += F::ONE;
                     }
-                    (flag_share, location_share)
+                    let t_after_compute_flag_loc = Instant::now();
+                    (
+                        flag_share,
+                        location_share,
+                        t_after_masked_address_tag,
+                        t_after_dpf_keygen,
+                        t_after_compute_flag_loc,
+                    )
                 }
             }
             _ => panic!("invalid party id"),
@@ -378,6 +346,8 @@ where
             SelectProtocol::select(comm, flag_share, location_share, access_counter_share)?
         };
 
+        let t_after_location_share = Instant::now();
+
         // 5. Reshare <flag> among all three parties
         let flag_share = match self.party_id {
             PARTY_1 => {
@@ -392,33 +362,127 @@ where
             _ => flag_share,
         };
 
+        let t_after_flag_share = Instant::now();
+
         // 6. Read the value <val> from the stash (if <flag>) or read a zero value
-        let value_share = stash_read_value::<C, F, SPDPF>(
-            comm,
-            self.access_counter,
-            location_share,
-            &self.stash_values_share,
-        )?;
+        let (
+            value_share,
+            t_after_convert_to_replicated,
+            t_after_masked_index,
+            t_after_dpf_key_distr,
+        ) = {
+            // a) convert the stash into replicated secret sharing
+            let fut_prev = comm.receive_previous::<Vec<F>>()?;
+            comm.send_next(self.stash_values_share.to_vec())?;
+            let stash_values_share_prev = fut_prev.get()?;
+
+            let t_after_convert_to_replicated = Instant::now();
+
+            // b) mask and reconstruct the stash index <loc>
+            let index_bits = (self.access_counter as f64).log2().ceil() as u32;
+            assert!(index_bits <= 16);
+            let bit_mask = ((1 << index_bits) - 1) as u16;
+            let (masked_loc, r_prev, r_next) =
+                MaskIndexProtocol::mask_index(comm, index_bits, location_share)?;
+
+            let t_after_masked_index = Instant::now();
+
+            // c) use DPFs to read the stash value
+            let fut_prev = comm.receive_previous::<SPDPF::Key>()?;
+            let fut_next = comm.receive_next::<SPDPF::Key>()?;
+            {
+                let (dpf_key_prev, dpf_key_next) =
+                    SPDPF::generate_keys(1 << index_bits, masked_loc as u64, F::ONE);
+                comm.send_previous(dpf_key_prev)?;
+                comm.send_next(dpf_key_next)?;
+            }
+            let dpf_key_prev = fut_prev.get()?;
+            let dpf_key_next = fut_next.get()?;
+            let t_after_dpf_key_distr = Instant::now();
+            let mut value_share = F::ZERO;
+            for j in 0..self.access_counter {
+                let index_prev = ((j as u16 + r_prev) & bit_mask) as u64;
+                let index_next = ((j as u16 + r_next) & bit_mask) as u64;
+                value_share +=
+                    SPDPF::evaluate_at(&dpf_key_prev, index_prev) * self.stash_values_share[j];
+                value_share +=
+                    SPDPF::evaluate_at(&dpf_key_next, index_next) * stash_values_share_prev[j];
+            }
+            (
+                value_share,
+                t_after_convert_to_replicated,
+                t_after_masked_index,
+                t_after_dpf_key_distr,
+            )
+        };
+
+        let t_after_dpf_eval = Instant::now();
+
+        let runtimes = runtimes.map(|mut r| {
+            r.record(
+                ProtocolStep::ReadMaskedAddressTag,
+                t_after_masked_address_tag - t_start,
+            );
+            r.record(
+                ProtocolStep::ReadDpfKeyGen,
+                t_after_dpf_keygen - t_after_masked_address_tag,
+            );
+            r.record(
+                ProtocolStep::ReadLookupFlagLocation,
+                t_after_compute_flag_loc - t_after_dpf_keygen,
+            );
+            r.record(
+                ProtocolStep::ReadComputeLocation,
+                t_after_location_share - t_after_compute_flag_loc,
+            );
+            r.record(
+                ProtocolStep::ReadReshareFlag,
+                t_after_flag_share - t_after_location_share,
+            );
+            r.record(
+                ProtocolStep::ReadConvertToReplicated,
+                t_after_convert_to_replicated - t_after_flag_share,
+            );
+            r.record(
+                ProtocolStep::ReadComputeMaskedIndex,
+                t_after_masked_index - t_after_convert_to_replicated,
+            );
+            r.record(
+                ProtocolStep::ReadDpfKeyDistribution,
+                t_after_dpf_key_distr - t_after_masked_index,
+            );
+            r.record(
+                ProtocolStep::ReadDpfEvaluations,
+                t_after_dpf_key_distr - t_after_dpf_eval,
+            );
+            r
+        });
 
         self.state = State::AwaitingWrite;
-        Ok(StashStateShare {
-            flag: flag_share,
-            location: location_share,
-            value: value_share,
-        })
+        Ok((
+            StashStateShare {
+                flag: flag_share,
+                location: location_share,
+                value: value_share,
+            },
+            runtimes,
+        ))
     }
 
-    fn write<C: AbstractCommunicator>(
+    pub fn write_with_runtimes<C: AbstractCommunicator>(
         &mut self,
         comm: &mut C,
         instruction: InstructionShare<F>,
         stash_state: StashStateShare<F>,
         db_address_share: F,
         db_value_share: F,
-    ) -> Result<(), Error> {
+        runtimes: Option<Runtimes>,
+    ) -> Result<Option<Runtimes>, Error> {
         assert_eq!(self.state, State::AwaitingWrite);
         assert!(self.access_counter < self.stash_size);
 
+        let t_start = Instant::now();
+
         // 1. Compute tag y := PRF(k, <db_adr>) such that P2, P3 obtain y.
         match self.party_id {
             PARTY_1 => {
@@ -452,27 +516,64 @@ where
             _ => panic!("invalid party id"),
         }
 
+        let t_after_address_tag = Instant::now();
+
         // 2. Insert new triple (<db_adr>, <db_val>, <db_val> into stash.
         self.stash_addresses_share.push(db_address_share);
         self.stash_values_share.push(db_value_share);
         self.stash_old_values_share.push(db_value_share);
 
+        let t_after_store_triple = Instant::now();
+
         // 3. Update stash
         let previous_value_share =
             SelectProtocol::select(comm, stash_state.flag, stash_state.value, db_value_share)?;
+        let t_after_select_previous_value = Instant::now();
         let value_share = SelectProtocol::select(
             comm,
             instruction.operation,
             instruction.value - previous_value_share,
             F::ZERO,
         )?;
-        stash_write_value::<C, F, SPDPF>(
-            comm,
-            self.access_counter,
-            stash_state.location,
-            value_share,
-            &mut self.stash_values_share,
-        )?;
+        let t_after_select_value = Instant::now();
+        let (t_after_masked_index, t_after_dpf_key_distr) = {
+            // a) mask and reconstruct the stash index <loc>
+            let index_bits = {
+                let bits = usize::BITS - self.access_counter.leading_zeros();
+                if bits > 0 {
+                    bits
+                } else {
+                    1
+                }
+            };
+            assert!(index_bits <= 16);
+            let bit_mask = ((1 << index_bits) - 1) as u16;
+            let (masked_loc, r_prev, r_next) =
+                MaskIndexProtocol::mask_index(comm, index_bits, stash_state.location)?;
+
+            let t_after_masked_index = Instant::now();
+
+            // b) use DPFs to read the stash value
+            let fut_prev = comm.receive_previous::<SPDPF::Key>()?;
+            let fut_next = comm.receive_next::<SPDPF::Key>()?;
+            {
+                let (dpf_key_prev, dpf_key_next) =
+                    SPDPF::generate_keys(1 << index_bits, masked_loc as u64, value_share);
+                comm.send_previous(dpf_key_prev)?;
+                comm.send_next(dpf_key_next)?;
+            }
+            let dpf_key_prev = fut_prev.get()?;
+            let dpf_key_next = fut_next.get()?;
+            let t_after_dpf_key_distr = Instant::now();
+            for j in 0..=self.access_counter {
+                let index_prev = ((j as u16).wrapping_add(r_prev) & bit_mask) as u64;
+                let index_next = ((j as u16).wrapping_add(r_next) & bit_mask) as u64;
+                self.stash_values_share[j] += SPDPF::evaluate_at(&dpf_key_prev, index_prev);
+                self.stash_values_share[j] += SPDPF::evaluate_at(&dpf_key_next, index_next);
+            }
+            (t_after_masked_index, t_after_dpf_key_distr)
+        };
+        let t_after_dpf_eval = Instant::now();
 
         self.access_counter += 1;
         self.state = if self.access_counter == self.stash_size {
@@ -480,7 +581,92 @@ where
         } else {
             State::AwaitingRead
         };
-        Ok(())
+
+        let runtimes = runtimes.map(|mut r| {
+            r.record(ProtocolStep::WriteAddressTag, t_after_address_tag - t_start);
+            r.record(
+                ProtocolStep::WriteStoreTriple,
+                t_after_store_triple - t_after_address_tag,
+            );
+            r.record(
+                ProtocolStep::WriteSelectPreviousValue,
+                t_after_select_previous_value - t_after_store_triple,
+            );
+            r.record(
+                ProtocolStep::WriteSelectValue,
+                t_after_select_value - t_after_select_previous_value,
+            );
+            r.record(
+                ProtocolStep::WriteComputeMaskedIndex,
+                t_after_masked_index - t_after_select_value,
+            );
+            r.record(
+                ProtocolStep::WriteDpfKeyDistribution,
+                t_after_dpf_key_distr - t_after_masked_index,
+            );
+            r.record(
+                ProtocolStep::WriteDpfEvaluations,
+                t_after_dpf_eval - t_after_dpf_key_distr,
+            );
+            r
+        });
+
+        Ok(runtimes)
+    }
+}
+
+impl<F, SPDPF> Stash<F> for StashProtocol<F, SPDPF>
+where
+    F: PrimeField + LegendreSymbol + Serializable,
+    SPDPF: SinglePointDpf<Value = F>,
+    SPDPF::Key: Serializable,
+{
+    fn get_party_id(&self) -> usize {
+        self.party_id
+    }
+
+    fn get_stash_size(&self) -> usize {
+        self.stash_size
+    }
+
+    fn get_access_counter(&self) -> usize {
+        self.access_counter
+    }
+
+    fn reset(&mut self) {
+        *self = Self::new(self.party_id, self.stash_size);
+    }
+
+    fn init<C: AbstractCommunicator>(&mut self, comm: &mut C) -> Result<(), Error> {
+        self.init_with_runtimes(comm, None).map(|_| ())
+    }
+
+    fn read<C: AbstractCommunicator>(
+        &mut self,
+        comm: &mut C,
+        instruction: InstructionShare<F>,
+    ) -> Result<StashStateShare<F>, Error> {
+        self.read_with_runtimes(comm, instruction, None)
+            .map(|x| x.0)
+    }
+
+    fn write<C: AbstractCommunicator>(
+        &mut self,
+        comm: &mut C,
+        instruction: InstructionShare<F>,
+        stash_state: StashStateShare<F>,
+        db_address_share: F,
+        db_value_share: F,
+    ) -> Result<(), Error> {
+        self.write_with_runtimes(
+            comm,
+            instruction,
+            stash_state,
+            db_address_share,
+            db_value_share,
+            None,
+        )
+        .map(|_| ())
     }
 
     fn get_stash_share(&self) -> (&[F], &[F], &[F]) {