circuit_context.h

// Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once

#include <algorithm>
#include <algorithm>
#include <memory>

#include "core/paddlefl_mpc/mpc_protocol/abstract_network.h"
#include "prng_utils.h"

namespace aby3 {

using AbstractNetwork = paddle::mpc::AbstractNetwork;

class CircuitContext {
public:
  CircuitContext(size_t party, std::shared_ptr<AbstractNetwork> network,
                 const block &seed = g_zero_block,
                 const block &seed2 = g_zero_block) {
    init(party, network, seed, seed2);
  }

  CircuitContext(const CircuitContext &other) = delete;

  CircuitContext &operator=(const CircuitContext &other) = delete;

  void init(size_t party, std::shared_ptr<AbstractNetwork> network, block seed,
            block seed2) {
    set_party(party);
    set_network(network);

    if (equals(seed, g_zero_block)) {
      seed = block_from_dev_urandom();
    }

    if (equals(seed2, g_zero_block)) {
      seed2 = block_from_dev_urandom();
    }
    set_random_seed(seed, 0);
    // seed2 is private
    set_random_seed(seed2, 2);

    // 3 for 3-party computation
    size_t party_pre = (this->party() - 1 + 3) % 3;
    size_t party_next = (this->party() + 1) % 3;

    if (party == 1) {
      block recv_seed = this->network()->template recv<block>(party_next);
      this->network()->template send(party_pre, seed);
      seed = recv_seed;
    } else {
      this->network()->template send(party_pre, seed);
      seed = this->network()->template recv<block>(party_next);
    }

    set_random_seed(seed, 1);
  }

  void set_party(size_t party) {
    if (party >= 3) {
      // exception handling
    }
    _party = party;
  }

  void set_network(std::shared_ptr<AbstractNetwork> network) {
    _network = network;
  }

  AbstractNetwork *network() { return _network.get(); }

  void set_random_seed(const block &seed, size_t idx) {
    if (idx >= 3) {
      // exception handling
    }
    _prng[idx].set_seed(seed);
  }

  size_t party() const { return _party; }

  size_t pre_party() const { return (_party + 3 - 1) % 3; }

  size_t next_party() const { return (_party + 1) % 3; }

  template <typename T> T gen_random(bool next) { return _prng[next].get<T>(); }

  template <typename T, template <typename> class Tensor>
  void gen_random(Tensor<T> &tensor, bool next) {
    std::for_each(
        tensor.data(), tensor.data() + tensor.numel(),
        [this, next](T &val) { val = this->template gen_random<T>(next); });
  }

  template <typename T> T gen_random_private() { return _prng[2].get<T>(); }

  template <typename T, template <typename> class Tensor>
  void gen_random_private(Tensor<T> &tensor) {
    std::for_each(
        tensor.data(), tensor.data() + tensor.numel(),
        [this](T &val) { val = this->template gen_random_private<T>(); });
  }

  template <typename T> T gen_zero_sharing_arithmetic() {
    return _prng[0].get<T>() - _prng[1].get<T>();
  }

  template <typename T, template <typename> class Tensor>
  void gen_zero_sharing_arithmetic(Tensor<T> &tensor) {
    std::for_each(tensor.data(), tensor.data() + tensor.numel(),
                  [this](T &val) {
                    val = this->template gen_zero_sharing_arithmetic<T>();
                  });
  }

  template <typename T> T gen_zero_sharing_boolean() {
    return _prng[0].get<T>() ^ _prng[1].get<T>();
  }

  template <typename T, template <typename> class Tensor>
  void gen_zero_sharing_boolean(Tensor<T> &tensor) {
    std::for_each(
        tensor.data(), tensor.data() + tensor.numel(),
        [this](T &val) { val = this->template gen_zero_sharing_boolean<T>(); });
  }

  template <typename T, template <typename> class Tensor>
  void ot(size_t sender, size_t receiver, size_t helper,
          const Tensor<T> *choice, const Tensor<T> *m[2], Tensor<T> *buffer[2],
          Tensor<T> *ret) {
    // TODO: check tensor shape equals
    const size_t numel = buffer[0]->numel();
    if (party() == sender) {
      bool common = helper == next_party();
      this->template gen_random(*buffer[0], common);
      this->template gen_random(*buffer[1], common);
      for (size_t i = 0; i < numel; ++i) {
        buffer[0]->data()[i] ^= m[0]->data()[i];
        buffer[1]->data()[i] ^= m[1]->data()[i];
      }
      network()->template send(receiver, *buffer[0]);
      network()->template send(receiver, *buffer[1]);

    } else if (party() == helper) {
      bool common = sender == next_party();

      this->template gen_random(*buffer[0], common);
      this->template gen_random(*buffer[1], common);

      for (size_t i = 0; i < numel; ++i) {
        // TODO: check if choice is one bit
        buffer[0]->data()[i] =
            choice->data()[i] ? buffer[1]->data()[i] : buffer[0]->data()[i];
      }
      network()->template send(receiver, *buffer[0]);
    } else if (party() == receiver) {
      network()->template recv(sender, *buffer[0]);
      network()->template recv(sender, *buffer[1]);
      network()->template recv(helper, *ret);
      size_t i = 0;
      std::for_each(ret->data(), ret->data() + numel,
                    [&buffer, &i, choice, ret](T &in) {
                      // TODO: check if choice is one bit
                      bool c = choice->data()[i];
                      in ^= buffer[c]->data()[i];
                      ++i;
                    });
    }
  }

private:
  size_t _party;

  std::shared_ptr<AbstractNetwork> _network;

  PseudorandomNumberGenerator _prng[3];
};

} // namespace aby3