Ctxt.h

/* Copyright (C) 2012-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_CTXT_H
#define HELIB_CTXT_H
 
#include <cfloat> // DBL_MAX
#include <helib/DoubleCRT.h>
#include <helib/apiAttributes.h>
 
namespace helib {
struct CKKS;
struct BGV;
 
template <typename Scheme>
class Ptxt;
 
class KeySwitch;
class PubKey;
class SecKey;
 
class SKHandle
{
  long powerOfS, powerOfX, secretKeyID;
 
public:
  friend class Ctxt;
 
  SKHandle(long newPowerOfS = 0, long newPowerOfX = 1, long newSecretKeyID = 0)
  {
    powerOfS = newPowerOfS;
    powerOfX = newPowerOfX;
    secretKeyID = newSecretKeyID;
  }
 
  void setBase(long newSecretKeyID = -1)
  {
    powerOfS = 1;
    powerOfX = 1;
    if (newSecretKeyID >= 0)
      secretKeyID = newSecretKeyID;
  }
 
  bool isBase(long ofKeyID = 0) const
  {
    // If ofKeyID<0, only check that this is base of some key,
    // otherwise check that this is base of the given key
    return powerOfS == 1 && powerOfX == 1 &&
           (ofKeyID < 0 || secretKeyID == ofKeyID);
  }
 
  void setOne(long newSecretKeyID = -1)
  {
    powerOfS = 0;
    powerOfX = 1;
    if (newSecretKeyID >= 0)
      secretKeyID = newSecretKeyID;
  }
 
  bool isOne() const { return powerOfS == 0; }
 
  bool operator==(const SKHandle& other) const
  {
    if (powerOfS == 0 && other.powerOfS == 0)
      return true;
    else
      return powerOfS == other.powerOfS && powerOfX == other.powerOfX &&
             secretKeyID == other.secretKeyID;
  }
 
  bool operator!=(const SKHandle& other) const { return !(*this == other); }
 
  /* access functions */
 
  long getPowerOfS() const { return powerOfS; }
  long getPowerOfX() const { return powerOfX; }
  long getSecretKeyID() const { return secretKeyID; }
 
  bool mul(const SKHandle& a, const SKHandle& b)
  {
    // If either of the inputs is one, the output equals to the other input
    if (a.isOne()) {
      *this = b;
      return (b.secretKeyID >= 0);
    }
    if (b.isOne()) {
      *this = a;
      return (a.secretKeyID >= 0);
    }
 
    if (a.secretKeyID == -1 || b.secretKeyID == -1) {
      secretKeyID = -1; // -1 will be used to indicate an "error state"
      return false;
    }
 
    if (a.secretKeyID != b.secretKeyID) {
      secretKeyID = -1;
      return false;
    }
 
    if (a.powerOfX != b.powerOfX) {
      secretKeyID = -1;
      return false;
    }
 
    secretKeyID = a.secretKeyID;
    powerOfX = a.powerOfX;
    powerOfS = a.powerOfS + b.powerOfS;
    return true;
  }
 
  friend std::istream& operator>>(std::istream& s, SKHandle& handle);
 
  // Raw IO
  void read(std::istream& str);
  void write(std::ostream& str) const;
};
 
inline std::ostream& operator<<(std::ostream& s, const SKHandle& handle)
{
  return s << "[" << handle.getPowerOfS() << " " << handle.getPowerOfX() << " "
           << handle.getSecretKeyID() << "]";
}
 
class CtxtPart : public DoubleCRT
{
public:
  SKHandle skHandle; // The secret-key polynomial corresponding to this part
 
  bool operator==(const CtxtPart& other) const;
  bool operator!=(const CtxtPart& other) const { return !(*this == other); }
 
  // Constructors
 
  CtxtPart(const Context& _context, const IndexSet& s) : DoubleCRT(_context, s)
  {
    skHandle.setOne();
  }
 
  CtxtPart(const Context& _context,
           const IndexSet& s,
           const SKHandle& otherHandle) :
      DoubleCRT(_context, s), skHandle(otherHandle)
  {}
 
  // Copy constructors from the base class
  explicit CtxtPart(const DoubleCRT& other) : DoubleCRT(other)
  {
    skHandle.setOne();
  }
 
  CtxtPart(const DoubleCRT& other, const SKHandle& otherHandle) :
      DoubleCRT(other), skHandle(otherHandle)
  {}
 
  void read(std::istream& str);
  void write(std::ostream& str) const;
};
 
std::istream& operator>>(std::istream& s, CtxtPart& p);
std::ostream& operator<<(std::ostream& s, const CtxtPart& p);
 
struct ZeroCtxtLike_type
{}; // used to select a constructor
 
const ZeroCtxtLike_type ZeroCtxtLike = ZeroCtxtLike_type();
 
class Ctxt
{
  friend class PubKey;
  friend class SecKey;
  friend class BasicAutomorphPrecon;
 
  const Context& context;      // points to the parameters of this FHE instance
  const PubKey& pubKey;        // points to the public encryption key;
  std::vector<CtxtPart> parts; // the ciphertext parts
  IndexSet primeSet; // the primes relative to which the parts are defined
  long ptxtSpace;    // plaintext space for this ciphertext (either p or p^r)
 
  // a high-probability bound on the the noise magnitude
  NTL::xdouble noiseBound;
 
  long intFactor; // an integer factor to divide by on decryption (for BGV)
  NTL::xdouble ratFactor; // rational factor to divide on decryption (for CKKS)
  NTL::xdouble ptxtMag;   // bound on the plaintext size (for CKKS)
 
  // Create a tensor product of c1,c2. It is assumed that *this,c1,c2
  // are defined relative to the same set of primes and plaintext space,
  // and that *this DOES NOT point to the same object as c1,c2
  void tensorProduct(const Ctxt& c1, const Ctxt& c2);
 
  // Add/subtract a ciphertext part to/from a ciphertext. These are private
  // methods, they cannot update the noiseBound so they must be called
  // from a procedure that will eventually update that estimate.
  Ctxt& operator-=(const CtxtPart& part)
  {
    subPart(part);
    return *this;
  }
 
  Ctxt& operator+=(const CtxtPart& part)
  {
    addPart(part);
    return *this;
  }
 
  // NOTE: the matchPrimeSets business in the following routines
  // is DEPRECATED.  The requirement is that the prime set of part
  // must contain the prime set of *this, unless *this is empty
  // (in which case *this takes on the prime set of part).
  // If not, an exception is raised.
  // Also, if matchPrimeSets == true and the prime set of *this does
  // not contain the prime set of part, an exception is also raised.
 
  // Procedural versions with additional parameter
  void subPart(const CtxtPart& part, bool matchPrimeSet = false)
  {
    subPart(part, part.skHandle, matchPrimeSet);
  }
 
  void addPart(const CtxtPart& part, bool matchPrimeSet = false)
  {
    addPart(part, part.skHandle, matchPrimeSet);
  }
 
  void subPart(const DoubleCRT& part,
               const SKHandle& handle,
               bool matchPrimeSet = false)
  {
    addPart(part, handle, matchPrimeSet, true);
  }
 
  void addPart(const DoubleCRT& part,
               const SKHandle& handle,
               bool matchPrimeSet = false,
               bool negative = false);
 
  // Takes as arguments a ciphertext-part p relative to s' and a key-switching
  // matrix W = W[s'->s], use W to switch p relative to (1,s), and add the
  // result to *this.
  void keySwitchPart(const CtxtPart& p, const KeySwitch& W);
 
  // internal procedure used in key-switching
  void keySwitchDigits(const KeySwitch& W, std::vector<DoubleCRT>& digits);
 
  long getPartIndexByHandle(const SKHandle& handle) const
  {
    for (size_t i = 0; i < parts.size(); i++)
      if (parts[i].skHandle == handle)
        return i;
    return -1;
  }
 
  // Sanity-check: Check that prime-set is "valid", i.e. that it
  // contains either all the special primes or none of them
  bool verifyPrimeSet() const;
 
  // A private assignment method that does not check equality of context or
  // public key, this is needed when we copy the pubEncrKey member between
  // different public keys.
  Ctxt& privateAssign(const Ctxt& other);
 
  // explicitly multiply intFactor by e, which should be
  // in the interval [0, ptxtSpace)
  void mulIntFactor(long e);
 
public:
  // Default copy-constructor
  Ctxt(const Ctxt& other) = default;
 
  // VJS-FIXME: this was really a messy design choice to not
  // have ciphertext cosntructors that specify prime sets.
  // The default value of ctxtPrimes is kind of pointless.
 
  //__attribute__((deprecated))
  explicit Ctxt(const PubKey& newPubKey, long newPtxtSpace = 0); // constructor
 
  //__attribute__((deprecated))
  Ctxt(ZeroCtxtLike_type, const Ctxt& ctxt);
  // constructs a zero ciphertext with same public key and
  // plaintext space as ctxt
 
  void DummyEncrypt(const NTL::ZZX& ptxt, double size = -1.0);
 
  Ctxt& operator=(const Ctxt& other)
  { // public assignment operator
    assertEq(&context,
             &other.context,
             "Cannot assign Ctxts with different context");
    assertEq(&pubKey,
             &other.pubKey,
             "Cannot assign Ctxts with different pubKey");
    return privateAssign(other);
  }
 
  bool operator==(const Ctxt& other) const { return equalsTo(other); }
  bool operator!=(const Ctxt& other) const { return !equalsTo(other); }
 
  // a procedural variant with an additional parameter
  bool equalsTo(const Ctxt& other, bool comparePkeys = true) const;
 
  // Encryption and decryption are done by the friends [Pub|Sec]Key
 
  void negate();
 
  // Add/subtract another ciphertext
  Ctxt& operator+=(const Ctxt& other)
  {
    addCtxt(other);
    return *this;
  }
 
  Ctxt& operator-=(const Ctxt& other)
  {
    addCtxt(other, true);
    return *this;
  }
 
  void addCtxt(const Ctxt& other, bool negative = false);
 
  // Multiply by another ciphertext
  void multLowLvl(const Ctxt& other, bool destructive = false);
  Ctxt& operator*=(const Ctxt& other)
  {
    multLowLvl(other);
    return *this;
  }
 
  void automorph(long k); // Apply automorphism F(X) -> F(X^k) (gcd(k,m)=1)
  Ctxt& operator>>=(long k)
  {
    automorph(k);
    return *this;
  }
 
  void complexConj(); // Complex conjugate, same as automorph(m-1)
 
  void smartAutomorph(long k);
  // Apply F(X)->F(X^k) followed by re-linearization. The automorphism is
  // possibly evaluated via a sequence of steps, to ensure that we can
  // re-linearize the result of every step.
 
  void frobeniusAutomorph(long j);
 
  // Operators acting between ciphertexts and plaintext objects
 
  // BGV case
  Ctxt& operator+=(const Ptxt<BGV>& other);
 
  Ctxt& operator-=(const Ptxt<BGV>& other);
 
  Ctxt& operator*=(const Ptxt<BGV>& other);
 
  // CKKS case
  Ctxt& operator+=(const Ptxt<CKKS>& other);
 
  Ctxt& operator-=(const Ptxt<CKKS>& other);
 
  Ctxt& operator*=(const Ptxt<CKKS>& other);
 
  Ctxt& operator*=(const NTL::ZZX& poly);
 
  Ctxt& operator*=(const long scalar);
 
  void addConstant(const DoubleCRT& dcrt, double size = -1.0);
  void addConstant(const NTL::ZZX& poly, double size = -1.0);
 
  template <typename Scheme>
  void addConstant(const Ptxt<Scheme>& ptxt)
  {
    addConstant(ptxt.getPolyRepr());
  }
 
  void addConstant(const NTL::ZZ& c);
  void addConstantCKKS(std::pair</*numerator=*/long, /*denominator=*/long>);
  void addConstantCKKS(double x)
  { // FIXME: not enough precision when x is large
    addConstantCKKS(
        rationalApprox(x, /*denomBound=*/1 << getContext().alMod.getR()));
  }
 
  void addConstantCKKS(const DoubleCRT& dcrt,
                       NTL::xdouble size = NTL::xdouble(-1.0),
                       NTL::xdouble factor = NTL::xdouble(-1.0));
 
  void addConstantCKKS(const NTL::ZZX& poly,
                       NTL::xdouble size = NTL::xdouble(-1.0),
                       NTL::xdouble factor = NTL::xdouble(-1.0));
 
  void addConstantCKKS(const std::vector<std::complex<double>>& ptxt);
 
  void addConstantCKKS(const Ptxt<CKKS>& ptxt);
  void addConstantCKKS(const NTL::ZZ& c);
 
  void multByConstant(const DoubleCRT& dcrt, double size = -1.0);
  void multByConstant(const NTL::ZZX& poly, double size = -1.0);
  void multByConstant(const zzX& poly, double size = -1.0);
  void multByConstant(const NTL::ZZ& c);
 
  template <typename Scheme>
  void multByConstant(const Ptxt<Scheme>& ptxt)
  {
    multByConstant(ptxt.getPolyRepr());
  }
 
  // [[deprecated]]
  void multByConstantCKKS(double x)
  {
    if (this->isEmpty())
      return;
 
    if (x == 1)
      return;
 
    if (x == 0) {
      clear();
      return;
    }
 
    double size = std::abs(x);
    ptxtMag *= size;
    ratFactor /= size;
    if (x < 0)
      this->negate();
  }
 
  void multByConstantCKKS(std::pair<long, long> num) // rational number
  {
    multByConstantCKKS(double(num.first) / num.second);
  }
 
  void multByConstantCKKS(const DoubleCRT& dcrt,
                          NTL::xdouble size = NTL::xdouble(-1.0),
                          NTL::xdouble factor = NTL::xdouble(-1.0),
                          double roundingErr = -1.0);
 
  void multByConstantCKKS(const NTL::ZZX& poly,
                          NTL::xdouble size = NTL::xdouble(-1.0),
                          NTL::xdouble factor = NTL::xdouble(-1.0),
                          double roundingErr = -1.0)
  {
    DoubleCRT dcrt(poly, context, primeSet);
    multByConstantCKKS(dcrt, size, factor, roundingErr);
  }
 
  // TODO: Ptxt: refactor into single function for getting the
  // std::vector<cx_double> repr of slots
  void multByConstantCKKS(const Ptxt<CKKS>& ptxt);
  void multByConstantCKKS(const std::vector<std::complex<double>>& ptxt);
 
  void xorConstant(const DoubleCRT& poly, UNUSED double size = -1.0)
  {
    DoubleCRT tmp = poly;
    tmp *= -2;
    tmp += 1; // tmp = 1-2*poly
    multByConstant(tmp);
    addConstant(poly);
  }
 
  void xorConstant(const NTL::ZZX& poly, double size = -1.0)
  {
    xorConstant(DoubleCRT(poly, context, primeSet), size);
  }
 
  void nxorConstant(const DoubleCRT& poly, UNUSED double size = -1.0)
  {
    DoubleCRT tmp = poly;
    tmp *= 2;
    tmp -= 1;                      // 2a-1
    addConstant(NTL::to_ZZX(-1L)); // b11
    multByConstant(tmp);           // (b-1)(2a-1)
    addConstant(poly);             // (b-1)(2a-1)+a = 1-a-b+2ab
  }
 
  void nxorConstant(const NTL::ZZX& poly, double size = -1.0)
  {
    nxorConstant(DoubleCRT(poly, context, primeSet), size);
  }
 
  void divideByP();
 
  void multByP(long e = 1)
  {
    long p2e = NTL::power_long(context.zMStar.getP(), e);
    ptxtSpace *= p2e;
    multByConstant(NTL::to_ZZ(p2e));
  }
 
  // For backward compatibility
  void divideBy2();
  void extractBits(std::vector<Ctxt>& bits, long nBits2extract = 0);
 
  // Higher-level multiply routines
  void multiplyBy(const Ctxt& other);
  void multiplyBy2(const Ctxt& other1, const Ctxt& other2);
  void square() { multiplyBy(*this); }
  void cube() { multiplyBy2(*this, *this); }
 
  void power(long e); // in polyEval.cpp
 
 
 
  void reducePtxtSpace(long newPtxtSpace);
 
  // This method can be used to increase the plaintext space, but the
  // high-order digits that you get this way are noise. Do not use it
  // unless you know what you are doing.
  void hackPtxtSpace(long newPtxtSpace) { ptxtSpace = newPtxtSpace; }
 
  void bumpNoiseBound(double factor) { noiseBound *= factor; }
 
  void reLinearize(long keyIdx = 0);
  // key-switch to (1,s_i), s_i is the base key with index keyIdx
 
  Ctxt& cleanUp();
  // relinearize, then reduce, then drop special primes and small primes.
 
  // void reduce() const;
 
  void blindCtxt(const NTL::ZZX& poly);
 
  NTL::xdouble modSwitchAddedNoiseBound() const;
 
  void modUpToSet(const IndexSet& s);
 
  void modDownToSet(const IndexSet& s);
 
  void bringToSet(const IndexSet& s);
 
  // Finding the "natural" state of a ciphertext
  double naturalSize() const;       
  IndexSet naturalPrimeSet() const; 
 
  void dropSmallAndSpecialPrimes();
 
  double capacity() const
  {
    if (noiseBound <= 1.0)
      return context.logOfProduct(getPrimeSet());
    else
      return context.logOfProduct(getPrimeSet()) - log(noiseBound);
  }
 
  // VJS-FIXME:
  // For CKKS, the "total noise" for a Ctxt is
  // ptxtMag * ratFactor + noiseBound.  We should define a function
  // totalNoiseBound() that returns this value for CKKS, and
  // noiseBound o/w.  We might also define a function netCapacity
  // (or something) that is defined in terms of totalNoiseBound().
 
  long bitCapacity() const { return long(capacity() / log(2.0)); }
 
  double logOfPrimeSet() const { return context.logOfProduct(getPrimeSet()); }
 
  double rawModSwitch(std::vector<NTL::ZZX>& zzParts, long toModulus) const;
 
  //  void computePowers(std::vector<Ctxt>& v, long nPowers) const;
 
  void evalPoly(const NTL::ZZX& poly);
 
 
  void clear()
  { // set as an empty ciphertext
    primeSet = context.ctxtPrimes;
    parts.clear();
    noiseBound = NTL::to_xdouble(0.0);
    // VJS-FIXME: we should also make sure the other fields
    // are set to their default values for a new, empty ctxt.
  }
 
  bool isEmpty() const { return (parts.size() == 0); }
 
  bool inCanonicalForm(long keyID = 0) const
  {
    if (parts.size() > 2)
      return false;
    if (parts.size() > 0 && !parts[0].skHandle.isOne())
      return false;
    if (parts.size() > 1 && !parts[1].skHandle.isBase(keyID))
      return false;
    return true;
  }
 
  bool isCorrect() const
  {
    NTL::ZZ q = context.productOfPrimes(primeSet);
    return NTL::to_xdouble(q) > noiseBound * 2;
  }
 
  const Context& getContext() const { return context; }
  const PubKey& getPubKey() const { return pubKey; }
  const IndexSet& getPrimeSet() const { return primeSet; }
  long getPtxtSpace() const { return ptxtSpace; }
  const NTL::xdouble& getNoiseBound() const { return noiseBound; }
  const NTL::xdouble& getRatFactor() const { return ratFactor; }
  const NTL::xdouble& getPtxtMag() const { return ptxtMag; }
  void setPtxtMag(const NTL::xdouble& z) { ptxtMag = z; }
  long getKeyID() const;
 
  bool isCKKS() const { return (getContext().alMod.getTag() == PA_cx_tag); }
 
  // Return r such that p^r = ptxtSpace
  long effectiveR() const
  {
    long p = context.zMStar.getP();
    for (long r = 1, p2r = p; r < NTL_SP_NBITS; r++, p2r *= p) {
      if (p2r == ptxtSpace)
        return r;
      if (p2r > ptxtSpace)
        throw RuntimeError("ctxt.ptxtSpace is not of the form p^r");
    }
    throw RuntimeError("ctxt.ptxtSpace is not of the form p^r");
    return 0; // just to keep the compiler happy
  }
 
  double log_of_ratio() const
  {
    double logNoise =
        (getNoiseBound() <= 0.0) ? -DBL_MAX : log(getNoiseBound());
    double logMod =
        empty(getPrimeSet()) ? -DBL_MAX : context.logOfProduct(getPrimeSet());
    return logNoise - logMod;
  }
  friend std::istream& operator>>(std::istream& str, Ctxt& ctxt);
  friend std::ostream& operator<<(std::ostream& str, const Ctxt& ctxt);
 
  // Raw IO
  void write(std::ostream& str) const;
  void read(std::istream& str);
 
  // scale up c1, c2 so they have the same ratFactor
  static void equalizeRationalFactors(Ctxt& c1, Ctxt& c2);
};
 
// set out=prod_{i=0}^{n-1} v[j], takes depth log n and n-1 products
// out could point to v[0], but having it pointing to any other v[i]
// will make the result unpredictable.
void totalProduct(Ctxt& out, const std::vector<Ctxt>& v);
 
void incrementalProduct(std::vector<Ctxt>& v);
 
void innerProduct(Ctxt& result,
                  const std::vector<Ctxt>& v1,
                  const std::vector<Ctxt>& v2);
inline Ctxt innerProduct(const std::vector<Ctxt>& v1,
                         const std::vector<Ctxt>& v2)
{
  Ctxt ret(v1[0].getPubKey());
  innerProduct(ret, v1, v2);
  return ret;
}
 
void innerProduct(Ctxt& result,
                  const std::vector<Ctxt>& v1,
                  const std::vector<DoubleCRT>& v2);
inline Ctxt innerProduct(const std::vector<Ctxt>& v1,
                         const std::vector<DoubleCRT>& v2)
{
  Ctxt ret(v1[0].getPubKey());
  innerProduct(ret, v1, v2);
  return ret;
}
 
void innerProduct(Ctxt& result,
                  const std::vector<Ctxt>& v1,
                  const std::vector<NTL::ZZX>& v2);
inline Ctxt innerProduct(const std::vector<Ctxt>& v1,
                         const std::vector<NTL::ZZX>& v2)
{
  Ctxt ret(v1[0].getPubKey());
  innerProduct(ret, v1, v2);
  return ret;
}
 
void CheckCtxt(const Ctxt& c, const char* label);
 
void extractDigits(std::vector<Ctxt>& digits, const Ctxt& c, long r = 0);
// implemented in extractDigits.cpp
 
inline void extractDigits(std::vector<Ctxt>& digits,
                          const Ctxt& c,
                          long r,
                          bool shortCut)
{
  if (shortCut)
    std::cerr << "extractDigits: the shortCut flag is disabled\n";
  extractDigits(digits, c, r);
}
 
inline void Ctxt::extractBits(std::vector<Ctxt>& bits, long nBits2extract)
{
  extractDigits(bits, *this, nBits2extract);
}
 
// @brief Extract the mod-p digits of a mod-p^{r+e} ciphertext.
 
// extendExtractDigits assumes that the slots of *this contains integers mod
// p^{r+e} i.e., that only the free terms are nonzero. (If that assumptions
// does not hold then the result will not be a valid ciphertext anymore.)
//
// It returns in the slots of digits[j] the j'th-lowest digits from the
// integers in the slots of the input. Namely, the i'th slot of digits[j]
// contains the j'th digit in the p-base expansion of the integer in the i'th
// slot of the *this.  The plaintext space of digits[j] is mod p^{e+r-j}.
 
void extendExtractDigits(std::vector<Ctxt>& digits,
                         const Ctxt& c,
                         long r,
                         long e);
// implemented in extractDigits.cpp
 
} // namespace helib
 
#endif // ifndef HELIB_CTXT_H