divviup / libprio-rs

Implementation of Prio in Rust.
Mozilla Public License 2.0
103 stars 30 forks source link

Use an Arc internally instead of cloning Prio3 measurement share #1094

Open divergentdave opened 2 months ago

divergentdave commented 2 months ago

We currently clone the measurement share inside of Prio3::prepare_init(), out of the input share, to save a copy in the prepare state. If the input share and prepare state structs both used an Arc smart pointer for this, then we could shuffle the ownership around without a copy. This could potentially be a very large allocation, depending on VDAF parameters. See divviup/janus#3285.

Both Prio3InputShare and Prio3PrepareState have all their fields private, so this could be done transparently. However, note that both use prio::vdaf::Share, which is public, and exposes that it contains a Vec<F>. I think the best way to implement this would be to create a new enum similar to Share, but with an Arc<Vec<F>> in one variant, and abandon use of Share for the measurement share fields internally. This would save us from big copies while not adding any allocation or atomic instruction overhead to the helper share case (i.e. from using Arc<Share<F, SEED_SIZE>>).

divergentdave commented 2 months ago

I ran into a further problem while trying to implement this. Flp::truncate() currently takes in a Vec<F>. This is convenient for Count and Histogram, as they can just pass ownership of the measurement share back out again. However, if we're using an Arc internally, passing exclusive ownership to this method is a problem. prepare_next() is the caller of truncate(), and it takes in ownership of a Prio3PrepareState, uses truncate() to swap a measurement share for an output share (which may be the same thing), and returns an OutputShare inside a PrepareTransition. We will need to change the Flp trait to either take in a slice or an Arc<Vec<F>> to avoid reintroducing the clone we're trying to remove.