Add rebatching API for int (#116)

Summary:
Pull Request resolved: #116

For UDP protocol, we need to add a new type of gate. Rebatching gate. This type of gate allows to break a batch of values into smaller batches or combine several batches into a larger one.

This diff adds the APIs for int type to batching/unbatching

Reviewed By: elliottlawrence

Differential Revision: D34914176

fbshipit-source-id: e8038f51cc8b7f5e25de9bb9198e931ef43c35e0

Author: Ruiyu Zhu

fbpcf/frontend/Int.h

@@ -204,6 +204,14 @@ class Int {
   typename Int<isSigned, width, true, schedulerId, usingBatch>::ExtractedInt
   extractIntShare() const;
 
+  Int<isSigned, width, isSecret, schedulerId, usingBatch> batchingWith(
+      const std::vector<
+          Int<isSigned, width, isSecret, schedulerId, usingBatch>>& others)
+      const;
+
+  std::vector<Int<isSigned, width, isSecret, schedulerId, usingBatch>>
+  unbatching(std::shared_ptr<std::vector<uint32_t>> unbatchingStrategy) const;
+
  private:
   template <typename T>
   std::vector<UnitIntType> convertTo64BitIntVector(

fbpcf/frontend/Int_impl.h

@@ -560,4 +560,48 @@ Int<isSigned, width, isSecret, schedulerId, usingBatch>::convertBitsToInt(
   }
 }
 
+template <
+    bool isSigned,
+    int8_t width,
+    bool isSecret,
+    int schedulerId,
+    bool usingBatch>
+Int<isSigned, width, isSecret, schedulerId, usingBatch>
+Int<isSigned, width, isSecret, schedulerId, usingBatch>::batchingWith(
+    const std::vector<Int<isSigned, width, isSecret, schedulerId, usingBatch>>&
+        others) const {
+  static_assert(usingBatch, "Only batch values needs to rebatch!");
+  Int<isSigned, width, isSecret, schedulerId, usingBatch> rst;
+  size_t batchSize = others.size();
+  std::vector<Bit<true, schedulerId, usingBatch>> bits(batchSize);
+  for (size_t i = 0; i < width; i++) {
+    for (size_t j = 0; j < batchSize; j++) {
+      bits[j] = others.at(j).data_.at(i);
+    }
+    rst.data_[i] = data_.at(i).batchingWith(bits);
+  }
+  return rst;
+}
+
+template <
+    bool isSigned,
+    int8_t width,
+    bool isSecret,
+    int schedulerId,
+    bool usingBatch>
+std::vector<Int<isSigned, width, isSecret, schedulerId, usingBatch>>
+Int<isSigned, width, isSecret, schedulerId, usingBatch>::unbatching(
+    std::shared_ptr<std::vector<uint32_t>> unbatchingStrategy) const {
+  static_assert(usingBatch, "Only batch values needs to rebatch!");
+  std::vector<Int<isSigned, width, isSecret, schedulerId, usingBatch>> rst(
+      unbatchingStrategy->size());
+  for (size_t i = 0; i < width; i++) {
+    auto bitVec = data_.at(i).unbatching(unbatchingStrategy);
+    for (size_t j = 0; j < unbatchingStrategy->size(); j++) {
+      rst.at(j).data_.at(i) = bitVec.at(j);
+    }
+  }
+  return rst;
+}
+
 } // namespace fbpcf::frontend

fbpcf/frontend/test/IntTest.cpp

@@ -1386,4 +1386,52 @@ TEST(IntTest, testSubscriptBatch) {
   }
 }
 
+TEST(IntTest, testRebatch) {
+  const int8_t width = 15;
+  scheduler::SchedulerKeeper<0>::setScheduler(
+      std::make_unique<scheduler::PlaintextScheduler>(
+          scheduler::WireKeeper::createWithUnorderedMap()));
+  using secSignedIntBatch = Integer<Secret<Batch<Signed<width>>>, 0>;
+
+  int partyId = 2;
+
+  int size1 = 5;
+  int size2 = 2;
+  int size3 = 3;
+
+  std::vector<int64_t> v1(size1, (int64_t(1) << (width - 1)) - 1);
+  std::vector<int64_t> v2(size2, -1 - v1[0]);
+  std::vector<int64_t> v3(size3, 3);
+
+  secSignedIntBatch int1(v1, partyId);
+  secSignedIntBatch int2(v2, partyId);
+  secSignedIntBatch int3(v3, partyId);
+
+  std::vector<int64_t> expectedV(size1 + size2 + size3);
+  for (size_t i = 0; i < size1; i++) {
+    expectedV[i] = v1.at(i);
+  }
+  for (size_t i = 0; i < size2; i++) {
+    expectedV[size1 + i] = v2.at(i);
+  }
+
+  for (size_t i = 0; i < size3; i++) {
+    expectedV[size1 + size2 + i] = v3.at(i);
+  }
+
+  auto int4 = int1.batchingWith({int2, int3});
+
+  testVectorEq(int4.openToParty(partyId).getValue(), expectedV);
+
+  auto int123 = int4.unbatching(
+      std::make_shared<std::vector<uint32_t>>(std::vector<uint32_t>(
+          {static_cast<unsigned int>(size1),
+           static_cast<unsigned int>(size2),
+           static_cast<unsigned int>(size3)})));
+
+  testVectorEq(int123.at(0).openToParty(partyId).getValue(), v1);
+  testVectorEq(int123.at(1).openToParty(partyId).getValue(), v2);
+  testVectorEq(int123.at(2).openToParty(partyId).getValue(), v3);
+}
+
 } // namespace fbpcf::frontend