|
@@ -0,0 +1,222 @@
|
|
|
|
+#include "heapsampler.hpp"
|
|
|
|
+
|
|
|
|
+// In each 64-bit RegAS in the heap, the top bit (of the reconstructed
|
|
|
|
+// value) is 0, the next 42 bits are the random tag, and the low 21 bits
|
|
|
|
+// are the element value.
|
|
|
|
+
|
|
|
|
+#define HEAPSAMPLE_TAG_BITS 42
|
|
|
|
+#define HEAPSAMPLE_ELT_BITS 21
|
|
|
|
+
|
|
|
|
+// Make the next random tag
|
|
|
|
+void HeapSampler::make_randtag(MPCTIO &tio, yield_t &yield)
|
|
|
|
+{
|
|
|
|
+ // Make a uniformly random HEAPSAMPLE_TAG_BITS-bit tag. This needs
|
|
|
|
+ // to be RegXS in order for the sum (XOR) of P0 and P1's independent
|
|
|
|
+ // values to be uniform.
|
|
|
|
+ RegXS tagx;
|
|
|
|
+ tagx.randomize(HEAPSAMPLE_TAG_BITS);
|
|
|
|
+ mpc_xs_to_as(tio, yield, randtag, tagx);
|
|
|
|
+}
|
|
|
|
+
|
|
|
|
+// Compute the heap size (the smallest power of two strictly greater
|
|
|
|
+// than k) needed to store k elements
|
|
|
|
+static size_t heapsize(size_t k)
|
|
|
|
+{
|
|
|
|
+ size_t ret = 1;
|
|
|
|
+ while (ret <= k) {
|
|
|
|
+ ret <<= 1;
|
|
|
|
+ }
|
|
|
|
+ return ret;
|
|
|
|
+}
|
|
|
|
+
|
|
|
|
+// Return a random bit that reconstructs to 1 with probability k/m
|
|
|
|
+static RegBS weighted_coin(MPCTIO &tio, yield_t &yield, size_t k,
|
|
|
|
+ size_t m)
|
|
|
|
+{
|
|
|
|
+ RegAS limit;
|
|
|
|
+ limit.ashare = size_t((__uint128_t(k)<<63)/m) * !tio.player();
|
|
|
|
+ RegXS randxs;
|
|
|
|
+ randxs.randomize(63);
|
|
|
|
+ RegAS randas;
|
|
|
|
+ mpc_xs_to_as(tio, yield, randas, randxs);
|
|
|
|
+ CDPF cdpf = tio.cdpf(yield);
|
|
|
|
+ auto[lt, eq, gt] = cdpf.compare(tio, yield, randas-limit, tio.aes_ops());
|
|
|
|
+
|
|
|
|
+ return lt;
|
|
|
|
+}
|
|
|
|
+
|
|
|
|
+// Constructor for a HeapSampler that samples k items from a stream
|
|
|
|
+// of abritrary and unknown size, using O(k) memory
|
|
|
|
+HeapSampler::HeapSampler(MPCTIO &tio, yield_t &yield, size_t k)
|
|
|
|
+ : k(k), m(0), heap(tio.player(), heapsize(k))
|
|
|
|
+{
|
|
|
|
+ run_coroutines(tio, [&tio, this](yield_t &yield) {
|
|
|
|
+ heap.init(tio, yield);
|
|
|
|
+ }, [&tio, this](yield_t &yield) {
|
|
|
|
+ make_randtag(tio, yield);
|
|
|
|
+ });
|
|
|
|
+}
|
|
|
|
+
|
|
|
|
+// An element has arrived
|
|
|
|
+void HeapSampler::ingest(MPCTIO &tio, yield_t &yield, RegAS elt)
|
|
|
|
+{
|
|
|
|
+ ++m;
|
|
|
|
+ RegAS tagged_elt = (randtag << HEAPSAMPLE_ELT_BITS) + elt;
|
|
|
|
+ RegAS elt_to_insert = tagged_elt;
|
|
|
|
+
|
|
|
|
+ if (m > k) {
|
|
|
|
+ RegAS extracted_elt;
|
|
|
|
+ RegBS selection_bit;
|
|
|
|
+ run_coroutines(tio, [&tio, this, &extracted_elt](yield_t &yield) {
|
|
|
|
+ extracted_elt = heap.extract_min(tio, yield);
|
|
|
|
+ }, [&tio, this, &selection_bit](yield_t &yield) {
|
|
|
|
+ selection_bit = weighted_coin(tio, yield, k, m);
|
|
|
|
+ });
|
|
|
|
+ mpc_select(tio, yield, elt_to_insert, selection_bit,
|
|
|
|
+ extracted_elt, tagged_elt);
|
|
|
|
+ }
|
|
|
|
+
|
|
|
|
+ run_coroutines(tio, [&tio, this, elt_to_insert](yield_t &yield) {
|
|
|
|
+ heap.insert_optimized(tio, yield, elt_to_insert);
|
|
|
|
+ }, [&tio, this](yield_t &yield) {
|
|
|
|
+ make_randtag(tio, yield);
|
|
|
|
+ });
|
|
|
|
+}
|
|
|
|
+
|
|
|
|
+// The stream has ended; output min(k,m) randomly sampled elements.
|
|
|
|
+// After calling this function, the HeapSampler is reset to its
|
|
|
|
+// initial m=0 state.
|
|
|
|
+std::vector<RegAS> HeapSampler::close(MPCTIO &tio, yield_t &yield)
|
|
|
|
+{
|
|
|
|
+ size_t retsize = k;
|
|
|
|
+ if (m < k) {
|
|
|
|
+ retsize = m;
|
|
|
|
+ }
|
|
|
|
+ std::vector<RegAS> ret(retsize);
|
|
|
|
+ for (size_t i=0; i<retsize; ++i) {
|
|
|
|
+ ret[i] = heap.extract_min(tio, yield);
|
|
|
|
+ ret[i] &= ((size_t(1)<<HEAPSAMPLE_ELT_BITS)-1);
|
|
|
|
+ }
|
|
|
|
+ // Compare each output to (size_t(1)<<HEAPSAMPLE_ELT_BITS), since
|
|
|
|
+ // there may be a carry; if the output is greater than or equal to
|
|
|
|
+ // that value, fix the carry
|
|
|
|
+ RegAS limit;
|
|
|
|
+ limit.ashare = (size_t(1)<<HEAPSAMPLE_ELT_BITS) * !tio.player();
|
|
|
|
+
|
|
|
|
+ std::vector<coro_t> coroutines;
|
|
|
|
+ for (size_t i=0; i<retsize; ++i) {
|
|
|
|
+ coroutines.emplace_back([&tio, &ret, i, limit](yield_t &yield) {
|
|
|
|
+ CDPF cdpf = tio.cdpf(yield);
|
|
|
|
+ auto[lt, eq, gt] = cdpf.compare(tio, yield,
|
|
|
|
+ ret[i]-limit, tio.aes_ops());
|
|
|
|
+ RegAS fix, zero;
|
|
|
|
+ mpc_select(tio, yield, fix, gt^eq, zero, limit);
|
|
|
|
+ ret[i] -= fix;
|
|
|
|
+ });
|
|
|
|
+ }
|
|
|
|
+ run_coroutines(tio, coroutines);
|
|
|
|
+ heap.init(tio, yield);
|
|
|
|
+ return ret;
|
|
|
|
+}
|
|
|
|
+
|
|
|
|
+void heapsampler_test(MPCIO &mpcio, const PRACOptions &opts, char **args)
|
|
|
|
+{
|
|
|
|
+ size_t n = 100;
|
|
|
|
+ size_t k = 10;
|
|
|
|
+
|
|
|
|
+ // The number of elements to stream
|
|
|
|
+ if (*args) {
|
|
|
|
+ n = atoi(*args);
|
|
|
|
+ ++args;
|
|
|
|
+ }
|
|
|
|
+ // The size of the random sample
|
|
|
|
+ if (*args) {
|
|
|
|
+ k = atoi(*args);
|
|
|
|
+ ++args;
|
|
|
|
+ }
|
|
|
|
+
|
|
|
|
+ MPCTIO tio(mpcio, 0, opts.num_threads);
|
|
|
|
+ run_coroutines(tio, [&mpcio, &tio, n, k] (yield_t &yield) {
|
|
|
|
+
|
|
|
|
+ std::cout << "\n===== STREAMING =====\n";
|
|
|
|
+
|
|
|
|
+ HeapSampler sampler(tio, yield, k);
|
|
|
|
+
|
|
|
|
+ for (size_t i=0; i<n; ++i) {
|
|
|
|
+ // For ease of checking, just have the elements be in a
|
|
|
|
+ // simple sequence
|
|
|
|
+ RegAS elt;
|
|
|
|
+ elt.ashare = (i+1) * (1 + 0xfff*tio.player());
|
|
|
|
+ sampler.ingest(tio, yield, elt);
|
|
|
|
+ }
|
|
|
|
+
|
|
|
|
+ std::vector<RegAS> sample = sampler.close(tio, yield);
|
|
|
|
+
|
|
|
|
+ tio.sync_lamport();
|
|
|
|
+ mpcio.dump_stats(std::cout);
|
|
|
|
+ mpcio.reset_stats();
|
|
|
|
+ tio.reset_lamport();
|
|
|
|
+
|
|
|
|
+ std::cout << "\n===== CHECKING =====\n";
|
|
|
|
+ size_t expected_size = k;
|
|
|
|
+ if (n < k) {
|
|
|
|
+ expected_size = n;
|
|
|
|
+ }
|
|
|
|
+ assert(sample.size() == expected_size);
|
|
|
|
+ std::vector<value_t> reconstructed_sample(expected_size);
|
|
|
|
+
|
|
|
|
+ std::vector<coro_t> coroutines;
|
|
|
|
+ for (size_t i=0; i<expected_size; ++i) {
|
|
|
|
+ coroutines.emplace_back(
|
|
|
|
+ [&tio, &sample, i, &reconstructed_sample](yield_t &yield) {
|
|
|
|
+ reconstructed_sample[i] = mpc_reconstruct(
|
|
|
|
+ tio, yield, sample[i]);
|
|
|
|
+ });
|
|
|
|
+ }
|
|
|
|
+ run_coroutines(tio, coroutines);
|
|
|
|
+ if (tio.player() == 0) {
|
|
|
|
+ for (size_t i=0; i<expected_size; ++i) {
|
|
|
|
+ printf("%06lx\n", reconstructed_sample[i]);
|
|
|
|
+ }
|
|
|
|
+ }
|
|
|
|
+ });
|
|
|
|
+}
|
|
|
|
+
|
|
|
|
+void weighted_coin_test(MPCIO &mpcio, const PRACOptions &opts, char **args)
|
|
|
|
+{
|
|
|
|
+ size_t iters = 100;
|
|
|
|
+ size_t m = 100;
|
|
|
|
+ size_t k = 10;
|
|
|
|
+
|
|
|
|
+ // The number of iterations
|
|
|
|
+ if (*args) {
|
|
|
|
+ iters = atoi(*args);
|
|
|
|
+ ++args;
|
|
|
|
+ }
|
|
|
|
+ // The denominator
|
|
|
|
+ if (*args) {
|
|
|
|
+ m = atoi(*args);
|
|
|
|
+ ++args;
|
|
|
|
+ }
|
|
|
|
+ // The numerator
|
|
|
|
+ if (*args) {
|
|
|
|
+ k = atoi(*args);
|
|
|
|
+ ++args;
|
|
|
|
+ }
|
|
|
|
+
|
|
|
|
+ MPCTIO tio(mpcio, 0, opts.num_threads);
|
|
|
|
+ run_coroutines(tio, [&mpcio, &tio, iters, m, k] (yield_t &yield) {
|
|
|
|
+
|
|
|
|
+ size_t heads = 0, tails = 0;
|
|
|
|
+ for (size_t i=0; i<iters; ++i) {
|
|
|
|
+ RegBS coin = weighted_coin(tio, yield, k, m);
|
|
|
|
+ bool coin_rec = mpc_reconstruct(tio, yield, coin);
|
|
|
|
+ if (coin_rec) {
|
|
|
|
+ heads++;
|
|
|
|
+ } else {
|
|
|
|
+ tails++;
|
|
|
|
+ }
|
|
|
|
+ printf("%lu flips %lu heads %lu tails\n", i+1, heads, tails);
|
|
|
|
+ }
|
|
|
|
+ });
|
|
|
|
+}
|