partialMatch.h

/* Copyright (C) 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_PARTIALMATCH_H
#define HELIB_PARTIALMATCH_H
 
#include <sstream>
#include <stack>
 
#include <helib/Matrix.h>
#include <helib/PolyMod.h>
 
// This code is in flux and should be considered bery alpha.
// Not recommended for public use.
 
namespace helib {
 
template <typename TXT>
inline Matrix<TXT> calculateMasks(const EncryptedArray& ea,
                                  Matrix<TXT> query,
                                  const Matrix<Ptxt<BGV>>& database)
{
  if (query.dims(0) != 1)
    throw InvalidArgument("Query must be a row vector");
  if (query.dims(1) != database.dims(1))
    throw InvalidArgument(
        "Database and query must have same number of columns");
  // TODO: case where query.dims(0) != database.dims(0)
 
  // Replicate the query once per row of the database
  // TODO: Some such replication will be needed once blocks/bands exist
  std::vector<long> columns(database.dims(0), 0l);
  Matrix<TXT>& mask = query;
  mask.transpose();
  mask = mask.columns(columns);
  mask.transpose();
 
  (mask -= database)
      .apply([&](auto& entry) { mapTo01(ea, entry); })
      .apply([](auto& entry) { entry.negate(); })
      .apply([](auto& entry) { entry.addConstant(NTL::ZZX(1l)); });
 
  return mask;
}
 
template <typename TXT>
Matrix<Ctxt> calculateMasks(const EncryptedArray& ea,
                            Matrix<TXT> query,
                            const Matrix<Ctxt>& database)
{
  if (query.dims(0) != 1)
    throw InvalidArgument("Query must be a row vector");
  if (query.dims(1) != database.dims(1))
    throw InvalidArgument(
        "Database and query must have same number of columns");
  // TODO: case where query.dims(0) != database.dims(0)
 
  // Replicate the query once per row of the database
  // TODO: Some such replication will be needed once blocks/bands exist
  std::vector<long> columns(database.dims(0), 0l);
  Matrix<TXT>& mask = query;
  mask.transpose();
  mask = mask.columns(columns);
  mask.transpose();
 
  (mask -= database)
      .apply([&](auto& entry) { mapTo01(ea, entry); })
      .apply([](auto& entry) { entry.negate(); })
      .apply([](auto& entry) { entry.addConstant(NTL::ZZX(1l)); });
 
  return mask;
}
 
template <typename TXT>
inline Matrix<TXT> calculateScores(
    const std::vector<std::vector<long>> index_sets,
    const std::vector<long>& offsets,
    const std::vector<Matrix<long>>& weights,
    const Matrix<TXT>& mask)
{
  assertEq<InvalidArgument>(index_sets.size(),
                            offsets.size(),
                            "index_sets and offsets must have matching size");
  assertEq<InvalidArgument>(index_sets.size(),
                            weights.size(),
                            "index_sets and weights must have matching size");
  auto ones(mask(0, 0));
  ones.clear();
  ones.addConstant(NTL::ZZX(1L));
  Matrix<TXT> result(ones, mask.dims(0), 1l);
  for (std::size_t i = 0; i < index_sets.size(); ++i) {
    const auto& index_set = index_sets.at(i);
    const auto& weight_set = weights.at(i);
    long offset = offsets.at(i);
 
    assertEq<InvalidArgument>(
        weight_set.dims(0),
        index_set.size(),
        "found mismatch between index set size and weight set size");
 
    assertEq<InvalidArgument>(weight_set.dims(1),
                              1lu,
                              "all weight sets must be column vectors");
 
    Matrix<TXT> submatrix = mask.columns(index_set);
    Matrix<TXT> factor(submatrix * weight_set);
    // factor should in fact be a 1*1 matrix
    factor.apply([&](auto& entry) { entry.addConstant(NTL::ZZX(offset)); });
    result.template entrywiseOperation<TXT>(
        factor,
        [](auto& lhs, const auto& rhs) -> decltype(auto) {
          lhs.multiplyBy(rhs);
          return lhs;
        });
  }
  return result;
}
 
inline PolyMod partialMatchEncode(uint32_t input, const Context& context)
{
  const long p = context.zMStar.getP();
  std::vector<long> coeffs(context.zMStar.getOrdP());
  // TODO - shouldn't keep checking input.
  for (long i = 0; i < long(coeffs.size()) && input != 0; ++i) {
    coeffs[i] = input % p;
    input /= p;
  }
  return PolyMod(coeffs, context.slotRing);
}
 
struct Expr;
class ColNumber;
 
using QueryExpr = std::shared_ptr<Expr>;
 
inline std::shared_ptr<ColNumber> makeQueryExpr(int cl)
{
  return std::make_shared<ColNumber>(cl);
}
 
struct Expr
{
  virtual std::string eval() const = 0;
  virtual ~Expr() = default;
};
 
class ColNumber : public Expr
{
public:
  std::string eval() const override { return std::to_string(column); }
 
  ColNumber(int c) : column(c) {}
 
private:
  int column;
};
 
class And : public Expr
{
public:
  std::string eval() const override
  {
    return lhs->eval() + " " + rhs->eval() + " &&";
  }
 
  And(const QueryExpr& l, const QueryExpr& r) : lhs(l), rhs(r) {}
 
private:
  QueryExpr lhs;
  QueryExpr rhs;
};
 
class Or : public Expr
{
public:
  std::string eval() const override
  {
    return lhs->eval() + " " + rhs->eval() + " ||";
  }
 
  Or(const QueryExpr& l, const QueryExpr& r) : lhs(l), rhs(r) {}
 
private:
  QueryExpr lhs;
  QueryExpr rhs;
};
 
inline std::shared_ptr<And> operator&&(const QueryExpr& lhs,
                                       const QueryExpr& rhs)
{
  return std::make_shared<And>(lhs, rhs);
}
 
inline std::shared_ptr<Or> operator||(const QueryExpr& lhs,
                                      const QueryExpr& rhs)
{
  return std::make_shared<Or>(lhs, rhs);
}
 
struct Query_t
{
  std::vector<std::vector<long>> Fs;
 
  std::vector<long> mus;
 
  std::vector<Matrix<long>> taus;
 
  bool containsOR = false;
 
  Query_t(const std::vector<std::vector<long>>& index_sets,
          const std::vector<long>& offsets,
          const std::vector<Matrix<long>>& weights,
          const bool isThereAnOR) :
      Fs(index_sets), mus(offsets), taus(weights), containsOR(isThereAnOR)
  {}
 
  Query_t(std::vector<std::vector<long>>&& index_sets,
          std::vector<long>&& offsets,
          std::vector<Matrix<long>>&& weights,
          bool isThereAnOR) :
      Fs(index_sets), mus(offsets), taus(weights), containsOR(isThereAnOR)
  {}
};
 
class QueryBuilder
{
 
  // 'outer' vec are the and groups and 'inner' are the or groups
  using vecvec = std::vector<std::vector<long>>;
 
public:
  QueryBuilder(const QueryExpr& expr) : query_str(expr->eval()) {}
 
  Query_t build(long columns) const
  {
 
    // Convert the query to "type 1" by expanding out necessary ORs
    vecvec expr = expandOr(query_str);
    bool containsOR = false;
 
    vecvec Fs(expr.size());
    {
      std::vector<long> v(columns);
      std::iota(v.begin(), v.end(), 0);
      std::fill(Fs.begin(), Fs.end(), v);
    }
    std::vector<long> mus(expr.size(), 1);
    std::vector<Matrix<long>> taus;
    taus.reserve(expr.size());
 
    // Create the taus
    for (long i = 0; i < long(expr.size()); ++i) { // Each tau
      mus[i] = 0;                                  // Set mu to zero.
      Matrix<long> M(columns, 1);                  // Create temp tau matrix
      containsOR = (expr[i].size() > 1) ? true : false;
      for (long j = 0; j < long(expr[i].size()); ++j) // Each column index
        M(expr[i][j], 0) = 1;                         // Mark those columns as 1
      taus.push_back(std::move(M));
    }
 
    return Query_t(std::move(Fs), std::move(mus), std::move(taus), containsOR);
  }
 
private:
  std::string query_str;
 
  void printStack(std::stack<vecvec> stack)
  {
    while (!stack.empty()) {
      printVecVec(stack.top());
      stack.pop();
    }
  }
 
  void printVecVec(const vecvec& vv)
  {
    for (const auto& v : vv) {
      std::cout << "[ ";
      for (const auto& e : v) {
        std::cout << e << " ";
      }
      std::cout << "]";
    }
    std::cout << "\n";
  }
 
  bool isNumber(const std::string& s) const
  {
    // Positive only
    return std::all_of(s.begin(), s.end(), ::isdigit);
  }
 
  vecvec expandOr(const std::string& s) const
  {
    std::stack<vecvec> convertStack;
 
    std::istringstream input{s};
    std::ostringstream output{};
 
    std::string symbol;
 
    while (input >> symbol) {
      if (!symbol.compare("&&")) {
        // Squash the top into penultimate.
        auto op = convertStack.top();
        convertStack.pop();
        auto& top = convertStack.top();
        top.insert(top.end(), op.begin(), op.end());
      } else if (!symbol.compare("||")) {
        // Cartesian-esque product
        auto op1 = convertStack.top();
        convertStack.pop();
        auto op2 = convertStack.top();
        convertStack.pop();
 
        vecvec prod;
        prod.reserve(op1.size() * op2.size());
        for (const auto& i : op1)
          for (const auto& j : op2) {
            auto x = i;
            x.insert(x.end(), j.begin(), j.end());
            prod.push_back(std::move(x));
          }
 
        convertStack.push(std::move(prod));
      } else {
        // Assume it is a number. But sanity check anyway.
        assertTrue(isNumber(symbol),
                   "String is not a number: '" + symbol + "'");
        convertStack.emplace(vecvec(1, {std::stol(symbol)}));
      }
    }
 
    // Now read answer off stack (should be size == 1).
    assertEq<LogicError>(1UL,
                         convertStack.size(),
                         "Size of stack after expandOr should be 1");
 
    return std::move(convertStack.top());
  }
};
 
template <typename TXT>
class Database
{
public:
  // FIXME: Generally, should Database own the Matrix uniquely?
  // Should we force good practice and ask that Context always be shared_ptr?
 
  // FIXME: Should probably move Matrix or make it unique_ptr or both?
  Database(const Matrix<TXT>& M, std::shared_ptr<const Context> c) :
      data(M), context(c)
  {}
 
  // FIXME: Should this option really exist?
  Database(const Matrix<TXT>& M, const Context& c) :
      data(M),
      context(std::shared_ptr<const helib::Context>(&c, [](auto UNUSED p) {}))
  {}
 
  // FIXME: Combination of TXT = ctxt and TXT2 = ptxt does not work
  template <typename TXT2>
  Matrix<TXT2> contains(const Query_t& lookup_query,
                        const Matrix<TXT2>& query_data) const;
 
  // FIXME: Combination of TXT = ctxt and TXT2 = ptxt does not work
  template <typename TXT2>
  Matrix<TXT2> getScore(const Query_t& weighted_query,
                        const Matrix<TXT2>& query_data) const;
 
  // TODO - correct name?
  long columns() { return data.dims(1); }
 
private:
  Matrix<TXT> data;
  std::shared_ptr<const Context> context;
};
 
template <typename TXT>
template <typename TXT2>
inline Matrix<TXT2> Database<TXT>::contains(
    const Query_t& lookup_query,
    const Matrix<TXT2>& query_data) const
{
  auto result = getScore<TXT2>(lookup_query, query_data);
 
  if (lookup_query.containsOR) {
    // FLT on the scores
    result.apply([&](auto& txt) {
      txt.power(context->alMod.getPPowR() - 1);
      return txt;
    });
  }
 
  return result;
}
 
template <typename TXT>
template <typename TXT2>
inline Matrix<TXT2> Database<TXT>::getScore(
    const Query_t& weighted_query,
    const Matrix<TXT2>& query_data) const
{
  auto mask = calculateMasks(*(context->ea), query_data, this->data);
 
  auto result = calculateScores(weighted_query.Fs,
                                weighted_query.mus,
                                weighted_query.taus,
                                mask);
 
  return result;
}
 
} // namespace helib
 
#endif