Ptxt.h

/* Copyright (C) 2019-2020 IBM Corp.
 * This program is 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. See accompanying LICENSE file.
 */
 
#ifndef HELIB_PTXT_H
#define HELIB_PTXT_H
 
#include <type_traits>
#include <vector>
#include <algorithm>
#include <numeric>
#include <iomanip>
 
#include <helib/Context.h>
#include <helib/EncryptedArray.h>
#include <helib/assertions.h>
#include <helib/PolyMod.h>
#include <helib/keys.h>
#include <helib/scheme.h>
#include <helib/EncodedPtxt.h>
 
namespace helib {
 
// Utility functions
 
template <typename From, typename Scheme>
inline std::vector<typename Scheme::SlotType> convertDataToSlotVector(
    const std::vector<From>& data,
    const Context& context)
{
  static_assert(std::is_same<Scheme, CKKS>::value ||
                    std::is_same<Scheme, BGV>::value,
                "Can only call convertDataToSlotVector with Scheme equals to "
                "CKKS or BGV");
  using To = typename Scheme::SlotType;
  std::vector<To> res(data.size(), Ptxt<Scheme>::convertToSlot(context, 0l));
  for (std::size_t i = 0; i < data.size(); ++i) {
    res[i] = data[i];
  }
  return res;
}
 
template <typename Scheme>
void innerProduct(Ptxt<Scheme>& result,
                  const std::vector<Ptxt<Scheme>>& first_vec,
                  const std::vector<Ptxt<Scheme>>& second_vec)
{
  static_assert(std::is_same<Scheme, CKKS>::value ||
                    std::is_same<Scheme, BGV>::value,
                "Can only call innerProduct with Scheme equals to CKKS or BGV");
  for (std::size_t i = 0; i < first_vec.size(); ++i) {
    assertTrue<RuntimeError>(first_vec[i].isValid(),
                             "Cannot call innerProduct on default-constructed"
                             " Ptxt as first argument at index " +
                                 std::to_string(i));
  }
  for (std::size_t i = 0; i < second_vec.size(); ++i) {
    assertTrue<RuntimeError>(second_vec[i].isValid(),
                             "Cannot call innerProduct on default-constructed"
                             " Ptxt as second argument at index " +
                                 std::to_string(i));
  }
  long n = std::min(first_vec.size(), second_vec.size());
  if (n <= 0) {
    result.clear();
    return;
  }
  result = first_vec[0];
  result *= second_vec[0];
  for (long i = 1; i < n; ++i)
    result += (first_vec[i] * second_vec[i]);
}
 
// Forward declaration as function is a friend of the templated `Ptxt` class.
template <typename Scheme>
std::istream& operator>>(std::istream& is, Ptxt<Scheme>& ptxt);
 
// Forward declaration as function is a friend of the templated `Ptxt` class.
template <typename Scheme>
std::ostream& operator<<(std::ostream& is, const Ptxt<Scheme>& ptxt);
 
template <typename Scheme>
class Ptxt
{
  static_assert(std::is_same<Scheme, CKKS>::value ||
                    std::is_same<Scheme, BGV>::value,
                "Can only create plaintext object parameterized by the crypto "
                "scheme (CKKS or BGV)");
 
public:
  static constexpr std::string_view typeName = "Ptxt";
 
  using SlotType = typename Scheme::SlotType;
 
  Ptxt();
 
  explicit Ptxt(const Context& context);
 
  explicit Ptxt(const PubKey& pk) : Ptxt(pk.getContext()){};
 
  Ptxt(const Context& context, const SlotType& value);
 
  template <typename U = Scheme,
            std::enable_if_t<std::is_same<U, BGV>::value>* = nullptr>
  Ptxt(const Context& context, const NTL::ZZX& value);
 
  Ptxt(const Context& context, const std::vector<SlotType>& data);
 
  template <typename T>
  Ptxt(const Context& context, const std::vector<T>& data) :
      Ptxt<Scheme>(context, convertDataToSlotVector<T, Scheme>(data, context))
  {}
 
  Ptxt(const Ptxt<Scheme>& other) = default;
 
  Ptxt(Ptxt<Scheme>&& other) noexcept = default;
 
  Ptxt<Scheme>& operator=(const Ptxt<Scheme>& v) = default;
 
  Ptxt<Scheme>& operator=(Ptxt<Scheme>&& v) noexcept = default;
 
  ~Ptxt() = default;
 
  bool isValid() const;
 
  size_t size() const;
 
  long lsize() const;
 
  const Context& getContext() const { return *context; }
 
  void setData(const std::vector<SlotType>& data);
 
  void setData(const SlotType& value);
 
  template <typename T = Scheme,
            typename std::enable_if_t<std::is_same<T, BGV>::value>* = nullptr>
  void setData(const NTL::ZZX& value);
 
  template <typename T = Scheme,
            typename std::enable_if_t<std::is_same<T, BGV>::value>* = nullptr>
  void decodeSetData(const NTL::ZZX& data);
 
  void clear();
 
  Ptxt<Scheme>& random();
 
  const std::vector<SlotType>& getSlotRepr() const;
 
  NTL::ZZX getPolyRepr() const;
 
  void encode(EncodedPtxt& eptxt,
              double mag = -1,
              OptLong prec = OptLong()) const;
 
  SlotType& operator[](long i);
 
  SlotType operator[](long i) const;
 
  SlotType& at(long i);
 
  SlotType at(long i) const;
 
  bool operator==(const Ptxt<Scheme>& other) const;
 
  bool operator!=(const Ptxt<Scheme>& other) const;
 
  Ptxt<Scheme> operator*(const Ptxt<Scheme>& rhs) const;
 
  Ptxt<Scheme> operator+(const Ptxt<Scheme>& rhs) const;
 
  Ptxt<Scheme> operator-(const Ptxt<Scheme>& rhs) const;
 
  Ptxt<Scheme>& operator*=(const Ptxt<Scheme>& otherPtxt);
 
  Ptxt<Scheme>& operator*=(const SlotType& scalar);
 
  template <typename Scalar>
  Ptxt<Scheme>& operator*=(const Scalar& scalar)
  {
    assertTrue<RuntimeError>(isValid(),
                             "Cannot call operator*= on "
                             "default-constructed Ptxt");
    for (std::size_t i = 0; i < this->slots.size(); i++) {
      this->slots[i] *= scalar;
    }
    return *this;
  }
 
  Ptxt<Scheme>& operator+=(const Ptxt<Scheme>& otherPtxt);
 
  Ptxt<Scheme>& operator+=(const SlotType& scalar);
 
  template <typename Scalar>
  Ptxt<Scheme>& operator+=(const Scalar& scalar)
  {
    assertTrue<RuntimeError>(isValid(),
                             "Cannot call operator+= on "
                             "default-constructed Ptxt");
    for (std::size_t i = 0; i < this->slots.size(); i++) {
      this->slots[i] += scalar;
    }
    return *this;
  }
 
  Ptxt<Scheme>& operator-=(const Ptxt<Scheme>& otherPtxt);
 
  Ptxt<Scheme>& operator-=(const SlotType& scalar);
 
  template <typename Scalar>
  Ptxt<Scheme>& operator-=(const Scalar& scalar)
  {
    assertTrue<RuntimeError>(isValid(),
                             "Cannot call operator-= on "
                             "default-constructed Ptxt");
    for (std::size_t i = 0; i < this->slots.size(); i++) {
      this->slots[i] -= scalar;
    }
    return *this;
  }
 
  Ptxt<Scheme>& negate();
 
  template <typename T = Scheme,
            typename Scalar,
            typename std::enable_if_t<std::is_same<T, BGV>::value>* = nullptr>
  Ptxt<Scheme>& addConstant(const Scalar& scalar)
  {
    return *this += scalar;
  }
 
  template <typename T = Scheme,
            typename Scalar,
            typename std::enable_if_t<std::is_same<T, CKKS>::value>* = nullptr>
  Ptxt<Scheme>& addConstantCKKS(const Scalar& scalar)
  {
    return *this += scalar;
  }
 
  Ptxt<Scheme>& multiplyBy(const Ptxt<Scheme>& otherPtxt);
 
  Ptxt<Scheme>& multiplyBy2(const Ptxt& otherPtxt1, const Ptxt& otherPtxt2);
 
  Ptxt<Scheme>& square();
 
  Ptxt<Scheme>& cube();
 
  Ptxt<Scheme>& power(long e);
 
  Ptxt<Scheme>& rotate(long amount);
 
  Ptxt<Scheme>& rotate1D(long dim, long amount);
 
  Ptxt<Scheme>& shift(long amount);
 
  Ptxt<Scheme>& shift1D(long dim, long amount);
 
  Ptxt<Scheme>& automorph(long k);
 
  template <typename T = Scheme,
            std::enable_if_t<std::is_same<T, BGV>::value>* = nullptr>
  Ptxt<Scheme>& frobeniusAutomorph(long j);
 
  Ptxt<Scheme>& replicate(long pos);
 
  std::vector<Ptxt<Scheme>> replicateAll() const;
 
  template <typename T = Scheme,
            std::enable_if_t<std::is_same<T, CKKS>::value>* = nullptr>
  Ptxt<Scheme>& complexConj();
 
  template <typename T = Scheme,
            std::enable_if_t<std::is_same<T, CKKS>::value>* = nullptr>
  Ptxt<Scheme> real() const;
 
  template <typename T = Scheme,
            std::enable_if_t<std::is_same<T, CKKS>::value>* = nullptr>
  Ptxt<Scheme> imag() const;
 
  Ptxt<Scheme>& runningSums();
 
  Ptxt<Scheme>& totalSums();
 
  Ptxt<Scheme>& incrementalProduct();
 
  Ptxt<Scheme>& totalProduct();
 
  Ptxt<Scheme>& mapTo01();
 
  void writeToJSON(std::ostream& os) const;
 
  JsonWrapper writeToJSON() const;
 
  static Ptxt<Scheme> readFromJSON(std::istream& is, const Context& context);
 
  static Ptxt<Scheme> readFromJSON(const JsonWrapper& jw,
                                   const Context& context);
 
  void readJSON(std::istream& is);
 
  void readJSON(const JsonWrapper& jw);
 
  friend std::istream& operator>><Scheme>(std::istream& is, Ptxt& ptxt);
 
  friend std::ostream& operator<<<Scheme>(std::ostream& os, const Ptxt& ptxt);
 
  static SlotType convertToSlot(const Context& context, long slot);
 
  Ptxt<Scheme>& cleanUp() { return *this; }
 
private:
  const Context* context;
 
  std::vector<SlotType> slots;
 
  long coordToIndex(const std::vector<long>& coords);
 
  std::vector<long> indexToCoord(long index);
 
  template <typename type>
  typename type::RX slotsToRX() const;
 
  void assertSlotsCompatible(const std::vector<SlotType>& slots) const;
 
  template <typename type>
  NTL::ZZX automorph_internal(long k);
};
} // namespace helib
 
#endif // HELIB_PTXT_H