helper funciton to split the encryption results. (#511)

Summary:
Pull Request resolved: #511

add a helper function to split encryption results into multiple shards

Reviewed By: chennyc

Differential Revision:
D43754424

Privacy Context Container: L416713

fbshipit-source-id: 61749ebc4d40f421285e7e4c5d629c95dddf3e37

fbpcf/mpc_std_lib/unified_data_process/data_processor/UdpUtil.cpp

@@ -11,6 +11,7 @@
 #include <boost/serialization/vector.hpp>
 #include <cstdint>
 #include <cstring>
+#include <iterator>
 #include "fbpcf/io/api/FileIOWrappers.h"
 
 namespace fbpcf::mpc_std_lib::unified_data_process::data_processor {
@@ -161,4 +162,33 @@ std::vector<__m128i> readExpandedKeyFromFile(const std::string& file) {
   return convertCharVecToM128i(data);
 }
 
+std::vector<IUdpEncryption::EncryptionResuts> splitEncryptionResults(
+    const IUdpEncryption::EncryptionResuts& encryptionResults,
+    int count) {
+  std::vector<IUdpEncryption::EncryptionResuts> rst;
+  rst.reserve(count);
+  size_t originalSize = encryptionResults.nonces.size();
+  auto ciphertextIt = encryptionResults.ciphertexts.begin();
+  auto noncesIt = encryptionResults.nonces.begin();
+  auto indexesIt = encryptionResults.indexes.begin();
+
+  for (size_t i = 0; i < count; i++) {
+    auto shardSize = getShardSize(originalSize, i, count);
+    rst.push_back(IUdpEncryption::EncryptionResuts{
+        .ciphertexts = std::vector<std::vector<unsigned char>>(
+            std::make_move_iterator(ciphertextIt),
+            std::make_move_iterator(ciphertextIt + shardSize)),
+        .nonces = std::vector<__m128i>(
+            std::make_move_iterator(noncesIt),
+            std::make_move_iterator(noncesIt + shardSize)),
+        .indexes = std::vector<int32_t>(
+            std::make_move_iterator(indexesIt),
+            std::make_move_iterator(indexesIt + shardSize))});
+    ciphertextIt += shardSize;
+    noncesIt += shardSize;
+    indexesIt += shardSize;
+  }
+  return rst;
+}
+
 } // namespace fbpcf::mpc_std_lib::unified_data_process::data_processor

fbpcf/mpc_std_lib/unified_data_process/data_processor/UdpUtil.h

@@ -62,6 +62,17 @@ std::vector<__m128i> readExpandedKeyFromFile(const std::string& file);
 IUdpEncryption::EncryptionResuts readEncryptionResultsFromFile(
     const std::string& file);
 
+std::vector<IUdpEncryption::EncryptionResuts> splitEncryptionResults(
+    const IUdpEncryption::EncryptionResuts& encryptionResults,
+    int count);
+
+// decides the i-th shard size. i is the shard index
+inline size_t
+getShardSize(size_t totalCount, size_t shardIndex, size_t totalShard) {
+  return (totalCount * (shardIndex + 1)) / totalShard -
+      (totalCount * shardIndex) / totalShard;
+}
+
 } // namespace fbpcf::mpc_std_lib::unified_data_process::data_processor
 
 #include "fbpcf/mpc_std_lib/unified_data_process/data_processor/UdpUtil_impl.h"

fbpcf/mpc_std_lib/unified_data_process/data_processor/test/DataProcessorTest.cpp

@@ -359,4 +359,38 @@ TEST(TestFileReadAndWrite, testEncryptionResultReadAndWrite) {
   }
 }
 
+TEST(TestSplit, getShardSize) {
+  EXPECT_EQ(getShardSize(11, 0, 3), 3);
+  EXPECT_EQ(getShardSize(11, 1, 3), 4);
+  EXPECT_EQ(getShardSize(11, 2, 3), 4);
+}
+
+TEST(TestSplit, testSplit) {
+  IUdpEncryption::EncryptionResuts results;
+  int batchSize = 11;
+  results.ciphertexts = std::vector<std::vector<unsigned char>>(batchSize);
+  results.indexes = std::vector<int32_t>(batchSize);
+  results.nonces = std::vector<__m128i>(batchSize);
+  for (size_t i = 0; i < batchSize; i++) {
+    results.nonces.at(i) = _mm_set_epi32(
+        folly::Random::rand32(),
+        folly::Random::rand32(),
+        folly::Random::rand32(),
+        folly::Random::rand32());
+    results.indexes.at(i) = folly::Random::rand32();
+    results.ciphertexts.at(i) = std::vector<unsigned char>(60);
+    for (size_t j = 0; j < results.ciphertexts.size(); j++) {
+      results.ciphertexts.at(i).at(j) = i + j;
+    }
+  }
+  auto resultsVec = splitEncryptionResults(results, 3);
+  for (auto& it : resultsVec) {
+    EXPECT_EQ(it.ciphertexts.size(), it.nonces.size());
+    EXPECT_EQ(it.ciphertexts.size(), it.indexes.size());
+  }
+  EXPECT_EQ(resultsVec.at(0).ciphertexts.size(), 3);
+  EXPECT_EQ(resultsVec.at(1).ciphertexts.size(), 4);
+  EXPECT_EQ(resultsVec.at(2).ciphertexts.size(), 4);
+}
+
 } // namespace fbpcf::mpc_std_lib::unified_data_process::data_processor