Add the ability to pre-expand RDPFs

rdpf.cpp

@@ -52,7 +52,7 @@ static value_t inverse_value_t(value_t x)
 // This algorithm is based on Appendix C from the Duoram paper, with a
 // small optimization noted below.
 RDPF::RDPF(MPCTIO &tio, yield_t &yield,
-    RegXS target, nbits_t depth)
+    RegXS target, nbits_t depth, bool save_expansion)
 {
     int player = tio.player();
     size_t &aesops = tio.aes_ops();
@@ -74,7 +74,12 @@ RDPF::RDPF(MPCTIO &tio, yield_t &yield,
     while(level < depth) {
         delete[] curlevel;
         curlevel = nextlevel;
-        nextlevel = new DPFnode[1<<(level+1)];
+        if (save_expansion && level == depth-1) {
+            expansion.resize(1<<depth);
+            nextlevel = expansion.data();
+        } else {
+            nextlevel = new DPFnode[1<<(level+1)];
+        }
         // Invariant: curlevel has 2^level elements; nextlevel has
         // 2^{level+1} elements
 
@@ -221,6 +226,10 @@ RDPF::RDPF(MPCTIO &tio, yield_t &yield,
                     bool flag = get_lsb(curlevel[i]);
                     DPFnode leftchild = xor_if(nextlevel[2*i], CW, flag);
                     DPFnode rightchild = xor_if(nextlevel[2*i+1], CWR, flag);
+                    if (save_expansion) {
+                        nextlevel[2*i] = leftchild;
+                        nextlevel[2*i+1] = rightchild;
+                    }
                     value_t leftlow = value_t(_mm_cvtsi128_si64x(leftchild));
                     value_t rightlow = value_t(_mm_cvtsi128_si64x(rightchild));
                     value_t lefthigh =
@@ -254,7 +263,9 @@ RDPF::RDPF(MPCTIO &tio, yield_t &yield,
     }
 
     delete[] curlevel;
-    delete[] nextlevel;
+    if (!save_expansion) {
+        delete[] nextlevel;
+    }
 }
 
 // The number of bytes it will take to store a RDPF of the given depth
@@ -293,6 +304,11 @@ DPFnode RDPF::descend(const DPFnode parent, nbits_t parentdepth,
 // Get the leaf node for the given input
 DPFnode RDPF::leaf(address_t input, size_t &op_counter) const
 {
+    // If we have a precomputed expansion, just use it
+    if (expansion.size()) {
+        return expansion[input];
+    }
+
     nbits_t totdepth = depth();
     DPFnode node = seed;
     for (nbits_t d=0;d<totdepth;++d) {
@@ -302,10 +318,52 @@ DPFnode RDPF::leaf(address_t input, size_t &op_counter) const
     return node;
 }
 
+// Expand the DPF if it's not already expanded
+void RDPF::expand(size_t &op_counter)
+{
+    nbits_t depth = this->depth();
+    size_t num_leaves = size_t(1)<<depth;
+    if (expansion.size() == num_leaves) return;
+    expansion.resize(num_leaves);
+    address_t index = 0;
+    address_t lastindex = 0;
+    DPFnode *path = new DPFnode[depth];
+    path[0] = seed;
+    for (nbits_t i=1;i<depth;++i) {
+        path[i] = descend(path[i-1], i-1, 0, op_counter);
+    }
+    expansion[index++] = descend(path[depth-1], depth-1, 0, op_counter);
+    expansion[index++] = descend(path[depth-1], depth-1, 1, op_counter);
+    while(index < num_leaves) {
+        // Invariant: lastindex and index will both be even, and
+        // index=lastindex+2
+        uint64_t index_xor = index ^ lastindex;
+        nbits_t how_many_1_bits = __builtin_popcount(index_xor);
+        // If lastindex -> index goes for example from (in binary)
+        // 010010110 -> 010011000, then index_xor will be
+        // 000001110 and how_many_1_bits will be 3.
+        // That indicates that path[depth-3] was a left child, and now
+        // we need to change it to a right child by descending right
+        // from path[depth-4], and then filling the path after that with
+        // left children.
+        path[depth-how_many_1_bits] =
+            descend(path[depth-how_many_1_bits-1],
+                depth-how_many_1_bits-1, 1, op_counter);
+        for (nbits_t i = depth-how_many_1_bits; i < depth-1; ++i) {
+            path[i+1] = descend(path[i], i, 0, op_counter);
+        }
+        lastindex = index;
+        expansion[index++] = descend(path[depth-1], depth-1, 0, op_counter);
+        expansion[index++] = descend(path[depth-1], depth-1, 1, op_counter);
+    }
+
+    delete[] path;
+}
+
 // Construct three RDPFs of the given depth all with the same randomly
 // generated target index.
 RDPFTriple::RDPFTriple(MPCTIO &tio, yield_t &yield,
-    nbits_t depth)
+    nbits_t depth, bool save_expansion)
 {
     // Pick a random XOR share of the target
     xs_target.randomize(depth);
@@ -316,7 +374,8 @@ RDPFTriple::RDPFTriple(MPCTIO &tio, yield_t &yield,
     for (int i=0;i<3;++i) {
         coroutines.emplace_back(
             [&, i](yield_t &yield) {
-                dpf[i] = RDPF(tio, yield, xs_target, depth);
+                dpf[i] = RDPF(tio, yield, xs_target, depth,
+                save_expansion);
             });
     }
     coroutines.emplace_back(

rdpf.hpp

@@ -28,6 +28,8 @@ struct RDPF {
     // XOR share of the scaling value M_xs such that the high words
     // of the leaf values for P0 and P1 XOR to M_xs * e_{target}
     RegXS scaled_xor;
+    // If we're saving the expansion, put it here
+    std::vector<DPFnode> expansion;
 
     RDPF() {}
 
@@ -44,7 +46,7 @@ struct RDPF {
     // 2^{depth+1}-2 local AES operations for P0,P1
     // 0 local AES operations for P2
     RDPF(MPCTIO &tio, yield_t &yield,
-        RegXS target, nbits_t depth);
+        RegXS target, nbits_t depth, bool save_expansion = false);
 
     // The number of bytes it will take to store this RDPF
     size_t size() const;
@@ -69,6 +71,9 @@ struct RDPF {
     // Cost: depth AES operations
     DPFnode leaf(address_t input, size_t &op_counter) const;
 
+    // Expand the DPF if it's not already expanded
+    void expand(size_t &op_counter);
+
     // Get the bit-shared unit vector entry from the leaf node
     inline RegBS unit_bs(DPFnode leaf) const {
         RegBS b;
@@ -121,6 +126,11 @@ T& operator>>(T &is, RDPF &rdpf)
     is.read((char *)&depth, sizeof(depth));
     rdpf.whichhalf = !!(depth & 0x80);
     depth &= 0x7f;
+    bool read_expanded = false;
+    if (depth > 64) {
+        read_expanded = true;
+        depth -= 64;
+    }
     assert(depth <= ADDRESS_MAX_BITS);
     rdpf.cw.clear();
     for (uint8_t i=0; i<depth; ++i) {
@@ -128,6 +138,11 @@ T& operator>>(T &is, RDPF &rdpf)
         is.read((char *)&cw, sizeof(cw));
         rdpf.cw.push_back(cw);
     }
+    if (read_expanded) {
+        rdpf.expansion.resize(1<<depth);
+        is.read((char *)rdpf.expansion.data(),
+            sizeof(rdpf.expansion[0])<<depth);
+    }
     value_t cfbits = 0;
     is.read((char *)&cfbits, BITBYTES(depth));
     rdpf.cfbits = cfbits;
@@ -138,20 +153,34 @@ T& operator>>(T &is, RDPF &rdpf)
     return is;
 }
 
+// Write the DPF to the output stream.  If expanded=true, then include
+// the expansion _if_ the DPF is itself already expanded.  You can use
+// this to write DPFs to files.
 template <typename T>
-T& operator<<(T &os, const RDPF &rdpf)
+T& write_maybe_expanded(T &os, const RDPF &rdpf,
+    bool expanded = true)
 {
     os.write((const char *)&rdpf.seed, sizeof(rdpf.seed));
     uint8_t depth = rdpf.cw.size();
     assert(depth <= ADDRESS_MAX_BITS);
     // The whichhalf bit is the high bit of depth
+    // If we're writing an expansion, add 64 to depth as well
     uint8_t whichhalf_and_depth = depth |
         (uint8_t(rdpf.whichhalf)<<7);
+    bool write_expansion = false;
+    if (expanded && rdpf.expansion.size() == (size_t(1)<<depth)) {
+        write_expansion = true;
+        whichhalf_and_depth += 64;
+    }
     os.write((const char *)&whichhalf_and_depth,
         sizeof(whichhalf_and_depth));
     for (uint8_t i=0; i<depth; ++i) {
         os.write((const char *)&rdpf.cw[i], sizeof(rdpf.cw[i]));
     }
+    if (write_expansion) {
+        os.write((const char *)rdpf.expansion.data(),
+            sizeof(rdpf.expansion[0])<<depth);
+    }
     os.write((const char *)&rdpf.cfbits, BITBYTES(depth));
     os.write((const char *)&rdpf.unit_sum_inverse, sizeof(rdpf.unit_sum_inverse));
     os.write((const char *)&rdpf.scaled_sum, sizeof(rdpf.scaled_sum));
@@ -160,6 +189,14 @@ T& operator<<(T &os, const RDPF &rdpf)
     return os;
 }
 
+// The ordinary << version never writes the expansion, since this is
+// what we use to send DPFs over the network.
+template <typename T>
+T& operator<<(T &os, const RDPF &rdpf)
+{
+    return write_maybe_expanded(os, rdpf, false);
+}
+
 // Computational peers will generate triples of RDPFs with the _same_
 // random target for use in Duoram.  They will each hold a share of the
 // target (neither knowing the complete target index).  They will each
@@ -179,15 +216,19 @@ struct RDPFTriple {
     // Construct three RDPFs of the given depth all with the same
     // randomly generated target index.
     RDPFTriple(MPCTIO &tio, yield_t &yield,
-        nbits_t depth);
+        nbits_t depth, bool save_expansion = false);
 };
 
 // I/O for RDPF Triples
 
+// We never write RDPFTriples over the network, so always write
+// the DPF expansions if they're available.
 template <typename T>
 T& operator<<(T &os, const RDPFTriple &rdpftrip)
 {
-    os << rdpftrip.dpf[0] << rdpftrip.dpf[1] << rdpftrip.dpf[2];
+    write_maybe_expanded(os, rdpftrip.dpf[0], true);
+    write_maybe_expanded(os, rdpftrip.dpf[1], true);
+    write_maybe_expanded(os, rdpftrip.dpf[2], true);
     nbits_t depth = rdpftrip.dpf[0].depth();
     os.write((const char *)&rdpftrip.as_target.ashare, BITBYTES(depth));
     os.write((const char *)&rdpftrip.xs_target.xshare, BITBYTES(depth));
@@ -212,10 +253,13 @@ struct RDPFPair {
 
 // I/O for RDPF Pairs
 
+// We never write RDPFPairs over the network, so always write
+// the DPF expansions if they're available.
 template <typename T>
 T& operator<<(T &os, const RDPFPair &rdpfpair)
 {
-    os << rdpfpair.dpf[0] << rdpfpair.dpf[1];
+    write_maybe_expanded(os, rdpfpair.dpf[0], true);
+    write_maybe_expanded(os, rdpfpair.dpf[1], true);
     return os;
 }