text:
/*******************************************************************************
 * Abstract domain described in Section 4 from the paper "An Abstract
 * Domain of Uninterpreted Functions" by Gange, Navas, Schachte,
 * Sondergaard, and Stuckey published in VMCAI'16.
 *
 * Each program variable is mapped to a syntactic term (aka
 * uninterpreted function). The join is antiunification and the meet
 * is a pseudo-meet based on the classical congruence closure
 * algorithm. The domain is suitable to infer equalities between
 * variables.
 ******************************************************************************/

#pragma once

#include <crab/domains/abstract_domain.hpp>
#include <crab/domains/abstract_domain_specialized_traits.hpp>
#include <crab/domains/backward_assign_operations.hpp>
#include <crab/domains/term/term_expr.hpp>
#include <crab/domains/term/term_operators.hpp>
#include <crab/support/debug.hpp>
#include <crab/support/stats.hpp>

#include <algorithm>
#include <map>
#include <set>
#include <utility>
#include <vector>

#include <boost/container/flat_map.hpp>
#include <boost/container/flat_set.hpp>
#include <boost/optional.hpp>

namespace crab {
namespace domains {

// TODO: factorize code. uf_domain and term_domain share a lot of
// code.
template <typename Number, typename VariableName>
class uf_domain final
    : public abstract_domain_api<uf_domain<Number, VariableName>> {

  using uf_domain_t = uf_domain<Number, VariableName>;
  using abstract_domain_t = abstract_domain_api<uf_domain_t>;

public:
  using typename abstract_domain_t::disjunctive_linear_constraint_system_t;
  using typename abstract_domain_t::interval_t;
  using typename abstract_domain_t::linear_constraint_system_t;
  using typename abstract_domain_t::linear_constraint_t;
  using typename abstract_domain_t::linear_expression_t;
  using typename abstract_domain_t::reference_constraint_t;
  using typename abstract_domain_t::variable_or_constant_t;
  using typename abstract_domain_t::variable_or_constant_vector_t;
  using typename abstract_domain_t::variable_t;
  using typename abstract_domain_t::variable_vector_t;
  using number_t = Number;
  using varname_t = VariableName;

private:
  using ttbl_t = term::term_table<number_t, term::term_operator_t>;
  using term_id_t = typename ttbl_t::term_id_t;
  using term_t = typename ttbl_t::term_t;
  using const_term_t = typename ttbl_t::const_term_t;
  using var_term_t = typename ttbl_t::var_term_t;
  using ftor_term_t = typename ttbl_t::ftor_term_t;

  using var_map_t = boost::container::flat_map<variable_t, term_id_t>;
  using var_set_t = std::set<variable_t>;
  using rev_var_map_t = boost::container::flat_map<term_id_t, var_set_t>;
  using linterm_t = typename linear_expression_t::component_t;

  bool m_is_bottom;
  ttbl_t m_ttbl;
  var_map_t m_var_map;
  rev_var_map_t m_rev_var_map;

  uf_domain(bool is_top) : m_is_bottom(!is_top) {}

  uf_domain(ttbl_t &&tbl, var_map_t &&vm, rev_var_map_t &&rvm)
      : m_is_bottom(false), m_ttbl(std::move(tbl)), m_var_map(std::move(vm)),
        m_rev_var_map(std::move(rvm)) {
    check_terms(__LINE__);
  }

  void check_terms(int line) const {
    CRAB_LOG(
        "uf-check-terms",
        for (auto const &p
             : m_var_map) {
          if (!(p.second < m_ttbl.size())) {
            CRAB_ERROR("term_equiv.hpp at line=", line, ": ",
                       "term id is not the table term");
          }
        }

        for (auto kv
             : m_rev_var_map) {
          for (auto v : kv.second) {
            auto it = m_var_map.find(v);
            if (it->second != kv.first) {
              CRAB_ERROR("term_equiv.hpp at line=", line, ": ", v,
                         " is mapped to t", it->second,
                         " but the reverse map says that should be t",
                         kv.first);
            }
          }
        });
  }

  void deref(term_id_t t) {
    std::vector<term_id_t> forgotten /*unused*/;
    m_ttbl.deref(t, forgotten);
  }

  /* Begin manipulate the reverse variable map */
  void add_rev_var_map(rev_var_map_t &rvmap, term_id_t t, variable_t v) const {
    auto it = rvmap.find(t);
    if (it != rvmap.end()) {
      it->second.insert(v);
    } else {
      var_set_t varset;
      varset.insert(v);
      rvmap.insert(std::make_pair(t, varset));
    }
  }

  void remove_rev_var_map(term_id_t t, const variable_t &v) {
    auto it = m_rev_var_map.find(t);
    if (it != m_rev_var_map.end()) {
      it->second.erase(v);
      if (it->second.empty()) {
        m_rev_var_map.erase(it);
      }
    }
  }
  /* End manipulate the reverse variable map */

  const term_t &get_term(term_id_t t) const {
    const term_t *ptr_t = m_ttbl.get_term_ptr(t);
    assert(ptr_t);
    return *ptr_t;
  }

  void rebind_var(const variable_t &x, term_id_t tx) {
    m_ttbl.add_ref(tx);

    auto it(m_var_map.find(x));
    if (it != m_var_map.end()) {
      remove_rev_var_map((*it).second, x);
      deref((*it).second);
      m_var_map.erase(it);
    }
    m_var_map.insert(std::make_pair(x, tx));
    add_rev_var_map(m_rev_var_map, tx, x);
  }

  term_id_t term_of_const(const number_t &n) {
    boost::optional<term_id_t> opt_tn(m_ttbl.find_const(n));
    if (opt_tn) {
      return *opt_tn;
    } else {
      return m_ttbl.make_const(n);
    }
  }

  term_id_t term_of_var(variable_t v, var_map_t &var_map,
                        rev_var_map_t &rvar_map, ttbl_t &ttbl) {
    auto it(var_map.find(v));
    if (it != var_map.end()) {
      // assert((*it).first == v);
      assert(ttbl.size() > (*it).second);
      return (*it).second;
    } else {
      // Allocate a fresh term
      term_id_t id(ttbl.fresh_var());
      var_map[v] = id;
      add_rev_var_map(rvar_map, id, v);
      ttbl.add_ref(id);
      return id;
    }
  }

  term_id_t term_of_var(variable_t v) {
    return term_of_var(v, m_var_map, m_rev_var_map, m_ttbl);
  }

  term_id_t term_of_linterm(linterm_t term) {
    if (term.first == 1) {
      return term_of_var(term.second);
    } else {
      return build_term(term::conv2termop(OP_MULTIPLICATION),
                        term_of_const(term.first), term_of_var(term.second));
    }
  }

  term_id_t build_term(term::term_operator_t op, term_id_t tx) {
    std::vector<term_id_t> ids = {tx};
    return build_term(op, ids);
  }

  term_id_t build_term(term::term_operator_t op, term_id_t tx, term_id_t ty) {
    std::vector<term_id_t> ids = {tx,ty};
    return build_term(op, ids);
  }

  term_id_t build_term(term::term_operator_t op, const std::vector<term_id_t> &ids) {
    boost::optional<term_id_t> eopt(m_ttbl.find_ftor(op, ids));
    if (eopt) {
      return *eopt;
    } else {
      term_id_t tx = m_ttbl.apply_ftor(op, ids);
      return tx;
    }
  }

  term_id_t build_linexpr(const linear_expression_t &e) {
    number_t cst = e.constant();
    typename linear_expression_t::const_iterator it(e.begin());
    if (it == e.end()) {
      return term_of_const(cst);
    }

    term_id_t t;
    if (cst == 0) {
      t = term_of_linterm(*it);
      ++it;
    } else {
      t = term_of_const(cst);
    }
    for (; it != e.end(); ++it) {
      t = build_term(term::conv2termop(OP_ADDITION), t, term_of_linterm(*it));
    }

    return t;
  }

  boost::optional<std::pair<variable_t, variable_t>>
  get_eq_or_diseq(linear_constraint_t cst) {
    if (cst.is_equality() || cst.is_disequation()) {
      if (cst.size() == 2 && cst.constant() == 0) {
        auto it = cst.begin();
        auto nx = it->first;
        auto vx = it->second;
        ++it;
        assert(it != cst.end());
        auto ny = it->first;
        auto vy = it->second;
        if (nx == (ny * -1)) {
          return std::make_pair(vx, vy);
        }
      }
    }
    return boost::optional<std::pair<variable_t, variable_t>>();
  }

  // helper for pseudo-meet: choose one non-var term from the
  // equivalence class associated with t.
  template <typename Range>
  boost::optional<term_id_t> choose_non_var(ttbl_t &ttbl,
                                            const Range &terms) const {
    std::vector<term_id_t> non_var_terms(terms.size());
    auto it = std::copy_if(terms.begin(), terms.end(), non_var_terms.begin(),
                           [&ttbl](term_id_t t) {
                             term_t *t_ptr = ttbl.get_term_ptr(t);
                             return (t_ptr && t_ptr->kind() == term::TERM_APP);
                           });
    non_var_terms.resize(std::distance(non_var_terms.begin(), it));
    if (non_var_terms.empty()) {
      return boost::optional<term_id_t>();
    } else {
      // TODO: the heuristics as described in the VMCAI'16 paper
      // that chooses the one that has more references each class.
      return *(non_var_terms.begin());
    }
  }

  // helper for pseudo-meet
  term_id_t build_dag_term(ttbl_t &ttbl, int t,
                           term::congruence_closure_solver<ttbl_t> &solver,
                           ttbl_t &out_ttbl, std::vector<int> &stack,
                           std::map<int, term_id_t> &cache) const {

    // already processed
    auto it = cache.find(t);
    if (it != cache.end()) {
      CRAB_LOG("uf-meet", crab::outs() << "build_dag_term. Found in cache: ";
               crab::outs() << "t" << t << " --> "
                            << "t" << it->second << "\n";);
      return it->second;
    }

    // break the cycle with a fresh variable
    if (std::find(stack.begin(), stack.end(), t) != stack.end()) {
      term_id_t v = out_ttbl.fresh_var();
      CRAB_LOG("uf-meet", crab::outs() << "build_dag_term. Detected cycle: ";
               crab::outs() << "t" << t << " --> "
                            << "t" << v << "\n";);
      return v;
    }

    stack.push_back(t);
    auto membs = solver.get_members(t);
    boost::optional<term_id_t> f = choose_non_var(ttbl, membs);

    if (!f) {
      // no concrete definition exists return a fresh variable
      term_id_t v = out_ttbl.fresh_var();
      auto res = cache.insert(std::make_pair(t, v));
      stack.pop_back();
      CRAB_LOG("uf-meet", crab::outs()
                              << "build_dag_term. No concrete definition: ";
               crab::outs() << "t" << t << " --> "
                            << "t" << (res.first)->second << "\n";);
      return (res.first)->second;
    } else {
      // traverse recursively the term
      term_t *f_ptr = ttbl.get_term_ptr(*f);
      CRAB_LOG("uf-meet",
               crab::outs()
                   << "build_dag_term. Traversing recursively the term "
                   << "t" << *f << ":";
               crab::outs() << *f_ptr << "\n";);
      const std::vector<term_id_t> &args(term::term_args(f_ptr));
      std::vector<term_id_t> res_args;
      res_args.reserve(args.size());
      for (term_id_t c : args) {
        res_args.push_back(build_dag_term(ttbl, solver.get_class(c), solver,
                                          out_ttbl, stack, cache));
      }
      auto res = cache.insert(std::make_pair(
          t, out_ttbl.apply_ftor(term::term_ftor(f_ptr), res_args)));
      stack.pop_back();
      CRAB_LOG("uf-meet", crab::outs()
                              << "build_dag_term. Finished recursive case: ";
               crab::outs() << "t" << t << " --> "
                            << "t" << (res.first)->second << "\n";);
      return (res.first)->second;
    }
  }

  void print_term(const term_t &t, crab_os &o) const {
    if (t.kind() == term::TERM_CONST) {
      const const_term_t *ct = static_cast<const const_term_t *>(&t);
      o << ct->val;
    } else if (t.kind() == term::TERM_VAR) {
      const var_term_t *vt = static_cast<const var_term_t *>(&t);
      o << "$VAR_" << vt->var;
    } else {
      assert(t.kind() == term::TERM_APP && "term should be a function");
      const ftor_term_t *ft = static_cast<const ftor_term_t *>(&t);
      o << ft->ftor << "(";
      for (unsigned i = 0, sz = ft->args.size(); i < sz;) {
        print_term(get_term(ft->args[i]), o);
        ++i;
        if (i < sz) {
          o << ",";
        }
      }
      o << ")";
    }
  }

public:
  uf_domain_t make_top() const override { return uf_domain_t(true); }

  uf_domain_t make_bottom() const override { return uf_domain_t(false); }

  void set_to_top() override {
    uf_domain abs(true);
    std::swap(*this, abs);
  }

  void set_to_bottom() override {
    uf_domain abs(false);
    std::swap(*this, abs);
  }

  uf_domain() : m_is_bottom(false) {}

  uf_domain(const uf_domain_t &o)
      : m_is_bottom(o.m_is_bottom), m_ttbl(o.m_ttbl), m_var_map(o.m_var_map),
        m_rev_var_map(o.m_rev_var_map) {
    crab::CrabStats::count(domain_name() + ".count.copy");
    crab::ScopedCrabStats __st__(domain_name() + ".copy");
    check_terms(__LINE__);
  }

  uf_domain_t &operator=(const uf_domain_t &o) {
    crab::CrabStats::count(domain_name() + ".count.copy");
    crab::ScopedCrabStats __st__(domain_name() + ".copy");

    o.check_terms(__LINE__);
    if (this != &o) {
      m_is_bottom = o.m_is_bottom;
      m_ttbl = o.m_ttbl;
      m_var_map = o.m_var_map;
      m_rev_var_map = o.m_rev_var_map;
    }
    check_terms(__LINE__);
    return *this;
  }

  bool is_bottom() const override { return m_is_bottom; }

  bool is_top() const override { return !m_var_map.size() && !is_bottom(); }

  // Lattice operations
  bool operator<=(const uf_domain_t &o) const override {
    crab::CrabStats::count(domain_name() + ".count.leq");
    crab::ScopedCrabStats __st__(domain_name() + ".leq");

    if (is_bottom()) {
      return true;
    } else if (o.is_bottom()) {
      return false;
    } else {
      // FIXME: avoid this copy
      uf_domain_t left(*this);
      uf_domain_t right(o);
      typename ttbl_t::term_map_t gen_map /*unused*/;

      // Build up the mapping of right onto left, variable by variable.
      // Assumption: the set of variables in left & right are common.
      for (auto p : left.m_var_map) {
        if (!left.m_ttbl.map_leq(right.m_ttbl, left.term_of_var(p.first),
                                 right.term_of_var(p.first), gen_map))
          return false;
      }
      return true;
    }
  }

  void operator|=(const uf_domain_t &o) override {
    crab::CrabStats::count(domain_name() + ".count.join");
    crab::ScopedCrabStats __st__(domain_name() + ".join");

    if (is_bottom() || o.is_top()) {
      *this = o;
    } else if (o.is_bottom() || is_top()) {
      return;
    } else {
      uf_domain_t right(o);
      ttbl_t out_tbl;
      var_map_t out_vmap;
      rev_var_map_t out_rvmap;
      typename ttbl_t::gener_map_t gener_map /*unused*/;

      for (auto p : m_var_map) {
        const variable_t &v = p.first;
        term_id_t tx = p.second;
        term_id_t ty(right.term_of_var(v));
        term_id_t tz =
            m_ttbl.generalize(right.m_ttbl, tx, ty, out_tbl, gener_map);
        assert(tz < out_tbl.size());
        out_vmap[v] = tz;
        add_rev_var_map(out_rvmap, tz, v);
      }

      for (auto p : out_vmap) {
        out_tbl.add_ref(p.second);
      }

      m_is_bottom = false;
      std::swap(m_ttbl, out_tbl);
      std::swap(m_var_map, out_vmap);
      std::swap(m_rev_var_map, out_rvmap);
    }
  }

  uf_domain_t operator|(const uf_domain_t &o) const override {
    crab::CrabStats::count(domain_name() + ".count.join");
    crab::ScopedCrabStats __st__(domain_name() + ".join");

    if (is_bottom() || o.is_top()) {
      return o;
    } else if (o.is_bottom() || is_top()) {
      return *this;
    } else {
      // FIXME: avoid this copy
      uf_domain_t left(*this);
      uf_domain_t right(o);
      ttbl_t out_tbl;
      typename ttbl_t::gener_map_t gener_map /*unused*/;
      var_map_t out_vmap;
      rev_var_map_t out_rvmap;

      // For each program variable in state, compute a generalization.
      for (auto p : left.m_var_map) {
        const variable_t &v = p.first;
        term_id_t tx = p.second;
        term_id_t ty = right.term_of_var(v);
        term_id_t tz =
            left.m_ttbl.generalize(right.m_ttbl, tx, ty, out_tbl, gener_map);
        assert(tz < out_tbl.size());
        out_vmap[v] = tz;
        add_rev_var_map(out_rvmap, tz, v);
      }

      for (auto p : out_vmap) {
        out_tbl.add_ref(p.second);
      }

      uf_domain_t res(std::move(out_tbl), std::move(out_vmap),
                      std::move(out_rvmap));

      CRAB_LOG("uf", crab::outs() << "============ JOIN ==================";
               crab::outs() << *this << "\n----------------";
               crab::outs() << o << "\n----------------";
               crab::outs() << res << "\n================"
                            << "\n");

      return res;
    }
  }

  uf_domain_t operator||(const uf_domain_t &other) const override {
    return *this | other;
  }

  uf_domain_t widening_thresholds(
      const uf_domain_t &other,
      const iterators::thresholds<number_t> &ts) const override {
    return *this | other;
  }

  // Meet
  uf_domain_t operator&(const uf_domain_t &o) const override {
    crab::CrabStats::count(domain_name() + ".count.meet");
    crab::ScopedCrabStats __st__(domain_name() + ".meet");

    if (is_bottom() || o.is_top()) {
      return *this;
    } else if (is_top() || o.is_bottom()) {
      return o;
    } else {
      ttbl_t out_ttbl(m_ttbl);
      std::map<term_id_t, term_id_t> copy_map;
      std::vector<int> stack;
      std::map<int, term_id_t> cache;
      var_map_t out_vmap;
      rev_var_map_t out_rvmap;

      // bring all terms to one ttbl
      for (auto p : o.m_var_map) {
        term_id_t tx = p.second;
        out_ttbl.copy_term(o.m_ttbl, tx, copy_map);
      }

      // build unifications between terms from this and o
      std::vector<std::pair<term_id_t, term_id_t>> eqs;
      for (auto p : m_var_map) {
        variable_t v(p.first);
        auto it = o.m_var_map.find(v);
        if (it != o.m_var_map.end()) {
          term_id_t tx = p.second;
          eqs.push_back(std::make_pair(tx, copy_map[it->second]));
        }
      }

      // compute equivalence classes
      term::congruence_closure_solver<ttbl_t> solver(&out_ttbl);
      solver.run(eqs);

      // new map from variable to an acyclic term
      for (auto p : m_var_map) {
        const variable_t &v = p.first;
        term_id_t t_old = p.second;
        term_id_t t_new = build_dag_term(out_ttbl, solver.get_class(t_old),
                                         solver, out_ttbl, stack, cache);
        out_vmap[v] = t_new;
        add_rev_var_map(out_rvmap, t_new, v);
      }
      for (auto p : o.m_var_map) {
        variable_t v(p.first);
        if (out_vmap.find(v) != out_vmap.end())
          continue;
        term_id_t t_old(copy_map[p.second]);
        term_id_t t_new = build_dag_term(out_ttbl, solver.get_class(t_old),
                                         solver, out_ttbl, stack, cache);
        out_vmap[v] = t_new;
        add_rev_var_map(out_rvmap, t_new, v);
      }

      for (auto p : out_vmap) {
        out_ttbl.add_ref(p.second);
      }

      uf_domain_t res(std::move(out_ttbl), std::move(out_vmap),
                      std::move(out_rvmap));

      CRAB_LOG("uf", crab::outs() << "============ MEET ==================";
               crab::outs() << *this << "\n----------------";
               crab::outs() << o << "\n----------------";
               crab::outs() << res << "\n================"
                            << "\n");
      return res;
    }
  }

  uf_domain_t operator&&(const uf_domain_t &o) const override {
    return *this & o;
  }

  // Remove a variable from the scope
  void operator-=(const variable_t &v) override {
    crab::CrabStats::count(domain_name() + ".count.forget");
    crab::ScopedCrabStats __st__(domain_name() + ".forget");

    auto it(m_var_map.find(v));
    if (it != m_var_map.end()) {
      term_id_t t = (*it).second;
      m_var_map.erase(it);
      remove_rev_var_map(t, v);
      deref(t);
    }
    CRAB_LOG("uf", crab::outs()
                       << "After removing " << v << ": " << *this << "\n";);
  }

  void assign(const variable_t &x, const linear_expression_t &e) override {
    crab::CrabStats::count(domain_name() + ".count.assign");
    crab::ScopedCrabStats __st__(domain_name() + ".assign");

    if (!is_bottom()) {
      term_id_t tx(build_linexpr(e));
      rebind_var(x, tx);
      check_terms(__LINE__);
      CRAB_LOG("uf", crab::outs() << "*** Assign " << x << ":=" << e << ":"
                                  << *this << "\n");
    }
  }

  // Apply operations to variables.

  // x = y op z
  void apply(arith_operation_t op, const variable_t &x, const variable_t &y,
             const variable_t &z) override {
    crab::CrabStats::count(domain_name() + ".count.apply");
    crab::ScopedCrabStats __st__(domain_name() + ".apply");
    check_terms(__LINE__);

    if (!is_bottom()) {
      term_id_t tx(
          build_term(term::conv2termop(op), term_of_var(y), term_of_var(z)));
      rebind_var(x, tx);
      check_terms(__LINE__);
      CRAB_LOG("uf", crab::outs() << "*** " << x << ":=" << y << " " << op
                                  << " " << z << ":" << *this << "\n");
    }
  }

  // x = y op k
  void apply(arith_operation_t op, const variable_t &x, const variable_t &y,
             number_t k) override {
    crab::CrabStats::count(domain_name() + ".count.apply");
    crab::ScopedCrabStats __st__(domain_name() + ".apply");

    if (!is_bottom()) {
      term_id_t tx(
          build_term(term::conv2termop(op), term_of_var(y), term_of_const(k)));
      rebind_var(x, tx);
      check_terms(__LINE__);
      CRAB_LOG("uf", crab::outs() << "*** " << x << ":=" << y << " " << op
                                  << " " << k << ":" << *this << "\n");
    }
  }

  void backward_assign(const variable_t &x, const linear_expression_t &e,
                       const uf_domain_t &inv) override {
    crab::CrabStats::count(domain_name() + ".count.backward_assign");
    crab::ScopedCrabStats __st__(domain_name() + ".backward_assign");
    if (!is_bottom()) {
      CRAB_WARN("backward_assign not implemented by ", domain_name());
    }
  }

  void backward_apply(arith_operation_t op, const variable_t &x,
                      const variable_t &y, number_t z,
                      const uf_domain_t &inv) override {
    crab::CrabStats::count(domain_name() + ".count.backward_apply");
    crab::ScopedCrabStats __st__(domain_name() + ".backward_apply");
    if (!is_bottom()) {
      CRAB_WARN("backward_apply not implemented by ", domain_name());
    }
  }

  void backward_apply(arith_operation_t op, const variable_t &x,
                      const variable_t &y, const variable_t &z,
                      const uf_domain_t &inv) override {
    crab::CrabStats::count(domain_name() + ".count.backward_apply");
    crab::ScopedCrabStats __st__(domain_name() + ".backward_apply");
    if (!is_bottom()) {
      CRAB_WARN("backward_apply not implemented by ", domain_name());
    }
  }

  void operator+=(const linear_constraint_t &cst) {
    crab::CrabStats::count(domain_name() + ".count.add_constraints");
    crab::ScopedCrabStats __st__(domain_name() + ".add_constraints");

    CRAB_LOG("uf", crab::outs()
                       << "*** Before assume " << cst << ":" << *this << "\n");

    if (is_bottom()) {
      return;
    }

    using pair_var_t = std::pair<variable_t, variable_t>;

    if (boost::optional<pair_var_t> eq = get_eq_or_diseq(cst)) {
      term_id_t tx(term_of_var((*eq).first));
      term_id_t ty(term_of_var((*eq).second));
      if (cst.is_disequation()) {
        if (tx == ty) {
          set_to_bottom();
          CRAB_LOG("uf", crab::outs() << "*** After assume " << cst << ":"
                                      << *this << "\n");
          return;
        }
      } else {
        // not bother if they are already equal
        if (tx == ty) {
          return;
        }

        std::vector<int> stack;
        std::map<int, term_id_t> cache;
        // congruence closure to compute equivalence classes
        term::congruence_closure_solver<ttbl_t> solver(&m_ttbl);
        std::vector<std::pair<term_id_t, term_id_t>> eqs = {{tx, ty}};
        solver.run(eqs);

        // new map from variable to an acyclic term
        for (auto p : m_var_map) {
          const variable_t &v = p.first;
          term_id_t t_old(term_of_var(v));
          term_id_t t_new = build_dag_term(m_ttbl, solver.get_class(t_old),
                                           solver, m_ttbl, stack, cache);
          rebind_var(v, t_new);
        }
      }
    }

    CRAB_LOG("uf", crab::outs()
                       << "*** After assume " << cst << ":" << *this << "\n");
  }

  void operator+=(const linear_constraint_system_t &csts) override {
    for (auto cst : csts) {
      this->operator+=(cst);
    }
  }

  interval_t operator[](const variable_t &x) override {
    crab::CrabStats::count(domain_name() + ".count.to_intervals");
    crab::ScopedCrabStats __st__(domain_name() + ".to_intervals");
    if (is_bottom()) {
      return interval_t::bottom();
    } else {
      return interval_t::top();
    }
  }

  void apply(int_conv_operation_t /*op*/, const variable_t &dst,
             const variable_t &src) override {
    // since reasoning about infinite precision we simply assign and
    // ignore the widths.
    assign(dst, src);
  }

  void apply(bitwise_operation_t op, const variable_t &x, const variable_t &y,
             const variable_t &z) override {
    crab::CrabStats::count(domain_name() + ".count.apply");
    crab::ScopedCrabStats __st__(domain_name() + ".apply");

    if (!is_bottom()) {
      term_id_t tx(
          build_term(term::conv2termop(op), term_of_var(y), term_of_var(z)));
      rebind_var(x, tx);
      check_terms(__LINE__);
      CRAB_LOG("uf", crab::outs() << "*** " << x << ":=" << y << " " << op
                                  << " " << z << ":" << *this << "\n");
    }
  }

  void apply(bitwise_operation_t op, const variable_t &x, const variable_t &y,
             number_t k) override {
    crab::CrabStats::count(domain_name() + ".count.apply");
    crab::ScopedCrabStats __st__(domain_name() + ".apply");

    if (!is_bottom()) {
      term_id_t tx(
          build_term(term::conv2termop(op), term_of_var(y), term_of_const(k)));
      rebind_var(x, tx);
      check_terms(__LINE__);
      CRAB_LOG("uf", crab::outs() << "*** " << x << ":=" << y << " " << op
                                  << " " << k << ":" << *this << "\n");
    }
  }

  /* Array operations */

  virtual void array_init(const variable_t & /*a*/,
                          const linear_expression_t & /*elem_size*/,
                          const linear_expression_t & /*lb_idx*/,
                          const linear_expression_t & /*ub_idx*/,
                          const linear_expression_t & /*val*/) override {
    // do nothing
  }

  virtual void array_load(const variable_t &lhs, const variable_t &a,
                          const linear_expression_t & /*elem_size*/,
                          const linear_expression_t &i) override {
    crab::CrabStats::count(domain_name() + ".count.array_read");
    crab::ScopedCrabStats __st__(domain_name() + ".array_read");

    if (!is_bottom()) {
      /**
       *  We treat the array load as an uninterpreted function
       *  lhs := array_load(a, i) -->  lhs := f(a,i)
       */
      term_id_t t_uf(
          build_term(term::TERM_OP_FUNCTION, term_of_var(a), build_linexpr(i)));
      rebind_var(lhs, t_uf);
      check_terms(__LINE__);
      CRAB_LOG("uf", crab::outs() << lhs << ":=" << a << "[" << i << "]  -- "
                                  << *this << "\n";);
    }
  }

  virtual void array_store(const variable_t &a,
                           const linear_expression_t & /*elem_size*/,
                           const linear_expression_t &i,
                           const linear_expression_t &val,
                           bool /*is_strong_update*/) override {
    // do nothing
  }

  virtual void array_store_range(const variable_t &a,
                                 const linear_expression_t &elem_size,
                                 const linear_expression_t &i,
                                 const linear_expression_t &j,
                                 const linear_expression_t &v) override {
    // do nothing
  }

  virtual void array_assign(const variable_t &lhs,
                            const variable_t &rhs) override {
    // do nothing
  }

  // backward array operations
  void backward_array_init(const variable_t &a,
                           const linear_expression_t &elem_size,
                           const linear_expression_t &lb_idx,
                           const linear_expression_t &ub_idx,
                           const linear_expression_t &val,
                           const uf_domain_t &invariant) override {
    CRAB_WARN(domain_name(), " does not implement backward operations");
  }
  void backward_array_load(const variable_t &lhs, const variable_t &a,
                           const linear_expression_t &elem_size,
                           const linear_expression_t &i,
                           const uf_domain_t &invariant) override {
    *this -= lhs;
    CRAB_WARN(domain_name(), " does not implement backward operations");
  }
  void backward_array_store(const variable_t &a,
                            const linear_expression_t &elem_size,
                            const linear_expression_t &i,
                            const linear_expression_t &v, bool is_strong_update,
                            const uf_domain_t &invariant) override {
    CRAB_WARN(domain_name(), " does not implement backward operations");
  }
  void backward_array_store_range(const variable_t &a,
                                  const linear_expression_t &elem_size,
                                  const linear_expression_t &i,
                                  const linear_expression_t &j,
                                  const linear_expression_t &v,
                                  const uf_domain_t &invariant) override {
    CRAB_WARN(domain_name(), " does not implement backward operations");
  }
  void backward_array_assign(const variable_t &lhs, const variable_t &rhs,
                             const uf_domain_t &invariant) override {
    CRAB_WARN(domain_name(), " does not implement backward operations");
  }

  DEFAULT_SELECT(uf_domain_t)

  // boolean operators
  virtual void assign_bool_cst(const variable_t &lhs,
                               const linear_constraint_t &rhs) override {
    // TODO
    operator-=(lhs);
  }

  virtual void assign_bool_ref_cst(const variable_t &lhs,
                                   const reference_constraint_t &rhs) override {
    // TODO
    operator-=(lhs);
  }

  virtual void assign_bool_var(const variable_t &lhs, const variable_t &rhs,
                               bool is_not_rhs) override {
    crab::CrabStats::count(domain_name() + ".count.assign_bool_var");
    crab::ScopedCrabStats __st__(domain_name() + ".assign_bool_var");

    if (!is_bottom()) {
      check_terms(__LINE__);
      if (is_not_rhs) {
        term_id_t tx(build_term(term::TERM_OP_NOT, term_of_var(rhs)));
        rebind_var(lhs, tx);
      } else {
        term_id_t tx(term_of_var(rhs));
        rebind_var(lhs, tx);
      }
      check_terms(__LINE__);

      CRAB_LOG(
          "uf", crab::outs() << "*** " << lhs << ":="; if (is_not_rhs) {
            crab::outs() << "not(" << rhs << ")";
          } else { crab::outs() << rhs; } crab::outs() << ":"
                                                       << *this << "\n");
    }
  }

  virtual void apply_binary_bool(bool_operation_t op, const variable_t &x,
                                 const variable_t &y,
                                 const variable_t &z) override {
    crab::CrabStats::count(domain_name() + ".count.apply_binary_bool");
    crab::ScopedCrabStats __st__(domain_name() + ".apply_binary_bool");

    if (!is_bottom()) {
      check_terms(__LINE__);
      term_id_t tx(
          build_term(term::conv2termop(op), term_of_var(y), term_of_var(z)));
      rebind_var(x, tx);
      check_terms(__LINE__);
      CRAB_LOG("uf", crab::outs() << "*** " << x << ":=" << y << " " << op
                                  << " " << z << ":" << *this << "\n");
    }
  }

  virtual void assume_bool(const variable_t &v, bool is_negated) override {
    // do nothing
  }

  void select_bool(const variable_t &lhs, const variable_t &cond,
                   const variable_t &b1, const variable_t &b2) override {
    operator-=(lhs);
  }

  // backward boolean operators
  virtual void backward_assign_bool_cst(const variable_t &lhs,
                                        const linear_constraint_t &rhs,
                                        const uf_domain_t &inv) override {
    CRAB_WARN(domain_name(), " does not implement backward operations");
  }

  virtual void backward_assign_bool_ref_cst(const variable_t &lhs,
                                            const reference_constraint_t &rhs,
                                            const uf_domain_t &inv) override {
    CRAB_WARN(domain_name(), " does not implement backward operations");
  }

  virtual void backward_assign_bool_var(const variable_t &lhs,
                                        const variable_t &rhs, bool is_not_rhs,
                                        const uf_domain_t &inv) override {
    CRAB_WARN(domain_name(), " does not implement backward operations");
  }

  virtual void backward_apply_binary_bool(bool_operation_t op,
                                          const variable_t &x,
                                          const variable_t &y,
                                          const variable_t &z,
                                          const uf_domain_t &inv) override {
    CRAB_WARN(domain_name(), " does not implement backward operations");
  }

  // Region operations
  virtual void region_init(const variable_t &reg) override {
    // do nothing
  }

  virtual void region_copy(const variable_t &lhs_reg,
                           const variable_t &rhs_reg) override {
    // do nothing
  }

  virtual void region_cast(const variable_t &src_reg,
                           const variable_t &dst_reg) override {
    // do nothing
  }

  virtual void ref_make(const variable_t &ref, const variable_t &reg,
                        const variable_or_constant_t &size,
                        const allocation_site &as) override {
    // do nothing
  }

  virtual void ref_free(const variable_t &reg, const variable_t &ref) override {
    // do nothing
  }

  virtual void ref_load(const variable_t &ref, const variable_t &reg,
                        const variable_t &res) override {
    crab::CrabStats::count(domain_name() + ".count.ref_load");
    crab::ScopedCrabStats __st__(domain_name() + ".ref_load");

    if (!is_bottom()) {
      /**
       *  We treat the load as an uninterpreted function:
       *  res := ref_load(reg, ref) -->  res := f(reg,ref)
       */
      term_id_t t_uf(build_term(term::TERM_OP_FUNCTION, term_of_var(reg),
                                build_linexpr(ref)));
      rebind_var(res, t_uf);
      check_terms(__LINE__);
      CRAB_LOG("uf", crab::outs() << res << ":=ref_load(" << reg << "," << ref
                                  << ")  -- " << *this << "\n";);
    }
  }

  virtual void ref_store(const variable_t &ref, const variable_t &reg,
                         const variable_or_constant_t &val) override {
    // do nothing
  }

  virtual void ref_gep(const variable_t &ref1, const variable_t &reg1,
                       const variable_t &ref2, const variable_t &reg2,
                       const linear_expression_t &offset) override {
    // do nothing
  }

  virtual void
  ref_load_from_array(const variable_t &lhs, const variable_t &ref,
                      const variable_t &region,
                      const linear_expression_t &index,
                      const linear_expression_t &elem_size) override {
    // TODO
    // do nothing
  }

  virtual void ref_store_to_array(const variable_t &ref,
                                  const variable_t &region,
                                  const linear_expression_t &index,
                                  const linear_expression_t &elem_size,
                                  const linear_expression_t &val) override {
    // do nothing
  }

  virtual void ref_assume(const reference_constraint_t &cst) override {
    // do nothing
  }

  void ref_to_int(const variable_t &reg, const variable_t &ref_var,
                  const variable_t &int_var) override {
    // do nothing
  }

  void int_to_ref(const variable_t &int_var, const variable_t &reg,
                  const variable_t &ref_var) override {
    // do nothing
  }

  void select_ref(const variable_t &lhs_ref, const variable_t &lhs_rgn,
                  const variable_t &cond, const variable_or_constant_t &ref1,
                  const boost::optional<variable_t> &rgn1,
                  const variable_or_constant_t &ref2,
                  const boost::optional<variable_t> &rgn2) override {
    // do nothing
  }

  boolean_value is_null_ref(const variable_t &ref) override {
    // do nothing
    return boolean_value();
  }
  bool
  get_allocation_sites(const variable_t &ref,
                       std::vector<allocation_site> &alloc_sites) override {
    // do nothing
    return false;
  }

  bool get_tags(const variable_t &rgn, const variable_t &ref,
                std::vector<uint64_t> &tags) override {
    // do nothing
    return false;
  }

  // Miscellaneous
  void rename(const variable_vector_t &from,
              const variable_vector_t &to) override {
    crab::CrabStats::count(domain_name() + ".count.rename");
    crab::ScopedCrabStats __st__(domain_name() + ".rename");

    if (is_top() || is_bottom()) {
      return;
    }

    CRAB_LOG(
        "uf", crab::outs() << "Renaming {"; for (auto v
                                                 : from) {
          crab::outs() << v << ";";
        } crab::outs() << "} with ";
        for (auto v
             : to) { crab::outs() << v << ";"; } crab::outs()
        << "}:\n";
        crab::outs() << *this << "\n";);

    auto error_if_found = [this](const variable_t &v) {
      auto it = m_var_map.find(v);
      if (it != m_var_map.end()) {
        CRAB_ERROR(domain_name() + "::rename assumes that ", v,
                   " does not exist");
      }
    };

    for (unsigned i = 0, sz = from.size(); i < sz; ++i) {
      const variable_t &v = from[i];
      const variable_t &new_v = to[i];
      if (v == new_v) { // nothing to rename
        continue;
      }

      error_if_found(new_v);

      auto it = m_var_map.find(v);
      if (it != m_var_map.end()) {
        term_id_t id = it->second;
        m_var_map.erase(it);
        m_var_map.insert(std::make_pair(new_v, id));
        remove_rev_var_map(id, v);
        add_rev_var_map(m_rev_var_map, id, new_v);
      }
    }
    CRAB_LOG("uf", crab::outs() << "RESULT=" << *this << "\n");
  }

  void forget(const variable_vector_t &variables) override {
    if (is_bottom() || is_top()) {
      return;
    }

    for (auto v : variables) {
      *this -= v;
    }
  }

  void project(const variable_vector_t &variables) override {
    crab::CrabStats::count(domain_name() + ".count.project");
    crab::ScopedCrabStats __st__(domain_name() + ".project");

    if (is_bottom() || is_top()) {
      return;
    }

    if (variables.empty()) {
      set_to_top();
      return;
    }

    std::set<variable_t> s1, s2;
    variable_vector_t s3;
    for (auto p : m_var_map) {
      s1.insert(p.first);
    }
    s2.insert(variables.begin(), variables.end());
    std::set_difference(s1.begin(), s1.end(), s2.begin(), s2.end(),
                        std::back_inserter(s3));
    forget(s3);
  }

  void expand(const variable_t &x, const variable_t &y) override {
    crab::CrabStats::count(domain_name() + ".count.expand");
    crab::ScopedCrabStats __st__(domain_name() + ".expand");

    if (is_bottom() || is_top()) {
      return;
    }

    linear_expression_t e(x);
    term_id_t tx(build_linexpr(e));
    rebind_var(y, tx);
    check_terms(__LINE__);
  }

  /* begin intrinsics operations */
  void intrinsic(std::string name, const variable_or_constant_vector_t &inputs,
                 const variable_vector_t &outputs) override {
    CRAB_WARN("Intrinsics ", name, " not implemented by ", domain_name());
  }

  void backward_intrinsic(std::string name,
                          const variable_or_constant_vector_t &inputs,
                          const variable_vector_t &outputs,
                          const uf_domain_t &invariant) override {
    CRAB_WARN("Intrinsics ", name, " not implemented by ", domain_name());
  }
  /* end intrinsics operations */

  void normalize() override {}
  void minimize() override {}

  // Output function
  void write(crab_os &o) const override {
    crab::CrabStats::count(domain_name() + ".count.write");
    crab::ScopedCrabStats __st__(domain_name() + ".write");

    if (is_bottom()) {
      o << "_|_";
      return;
    }
    if (m_var_map.empty()) {
      o << "{}";
      return;
    }

    bool first = true;
    o << "{";
    for (auto p : m_var_map) {
      if (first) {
        first = false;
      } else {
        o << ", ";
      }
      o << p.first << " -> ";
      print_term(get_term(p.second), o);
    }
    o << "}";

    CRAB_LOG("ufo-print-ttbl",
             /// For debugging purposes
             o << " ttbl={" << m_ttbl << "}\n";);
  }

  linear_constraint_system_t to_linear_constraint_system() const override {
    crab::CrabStats::count(domain_name() +
                           ".count.to_linear_constraint_system");
    crab::ScopedCrabStats __st__(domain_name() +
                                 ".to_linear_constraint_system");

    linear_constraint_system_t out_csts;
    if (is_bottom()) {
      out_csts += linear_constraint_t::get_false();
    } else if (!is_top()) {
      // Extract equalities

      // Seen equalities to avoid adding twice the same.
      std::set<std::pair<variable_t, variable_t>> seen;
      for (auto &kv : m_var_map) {
        const variable_t &x = kv.first;
        term_id_t tx = kv.second;
        auto it = m_rev_var_map.find(tx);
        if (it == m_rev_var_map.end()) {
          // this shouldn't happen
          continue;
        }
        for (auto var : it->second) {
          if (var.index() != x.index()) {
            if (seen.count(std::make_pair(var, x)) <= 0) {
              seen.insert(std::make_pair(x, var));
              out_csts += linear_constraint_t(linear_expression_t(x) ==
                                              linear_expression_t(var));
            }
          }
        }
      }
    }
    return out_csts;
  }

  disjunctive_linear_constraint_system_t
  to_disjunctive_linear_constraint_system() const override {
    auto lin_csts = to_linear_constraint_system();
    if (lin_csts.is_false()) {
      return disjunctive_linear_constraint_system_t(true /*is_false*/);
    } else if (lin_csts.is_true()) {
      return disjunctive_linear_constraint_system_t(false /*is_false*/);
    } else {
      return disjunctive_linear_constraint_system_t(lin_csts);
    }
  }

  std::string domain_name() const override { return "UFDomain"; }
}; // class uf_domain

template <typename Number, typename VariableName>
struct abstract_domain_traits<uf_domain<Number, VariableName>> {
  using number_t = Number;
  using varname_t = VariableName;
}; // end uf_domain

} // namespace domains
} // namespace crab

max_stars_repo_name: seahorn/crab

max_issues_repo_name: seahorn/crab

max_forks_repo_name: seahorn/crab


text:
/**
 * $Id$
 *
 * Copyright (C)
 * 2014 - $Date$
 *     Martin Wolf <ndhist@martin-wolf.org>
 *
 * This file is distributed under the BSD 2-Clause Open Source License
 * (See LICENSE file).
 *
 */
#include <boost/python.hpp>

#include <boost/numpy.hpp>

namespace ndhist {

void register_error_types();
void register_axis();
void register_ndhist();
void register_ndtable();
void register_stats_module();

namespace axes {

void register_generic_axis();
void register_linear_axis();
void register_log10_axis();

}//namespace axes

}//namespace ndhist

BOOST_PYTHON_MODULE(core)
{
    boost::numpy::initialize();

    ndhist::register_error_types();
    ndhist::register_axis();
    ndhist::axes::register_generic_axis();
    ndhist::axes::register_linear_axis();
    ndhist::axes::register_log10_axis();
    ndhist::register_ndhist();
    ndhist::register_ndtable();
    ndhist::register_stats_module();
}

ext: cpp

max_stars_repo_name: martwo/ndhist

max_issues_repo_name: martwo/ndhist

max_forks_repo_name: martwo/ndhist


text:
/**
*    Copyright (C) 2008 10gen Inc.
*
*    This program is free software: you can redistribute it and/or  modify
*    it under the terms of the GNU Affero General Public License, version 3,
*    as published by the Free Software Foundation.
*
*    This program is distributed in the hope that it will be useful,
*    but WITHOUT ANY WARRANTY; without even the implied warranty of
*    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
*    GNU Affero General Public License for more details.
*
*    You should have received a copy of the GNU Affero General Public License
*    along with this program.  If not, see <http://www.gnu.org/licenses/>.
*
*    As a special exception, the copyright holders give permission to link the
*    code of portions of this program with the OpenSSL library under certain
*    conditions as described in each individual source file and distribute
*    linked combinations including the program with the OpenSSL library. You
*    must comply with the GNU Affero General Public License in all respects
*    for all of the code used other than as permitted herein. If you modify
*    file(s) with this exception, you may extend this exception to your
*    version of the file(s), but you are not obligated to do so. If you do not
*    wish to do so, delete this exception statement from your version. If you
*    delete this exception statement from all source files in the program,
*    then also delete it in the license file.
*/

#include "mongo/pch.h"

#include <boost/filesystem/convenience.hpp>
#include <boost/filesystem/operations.hpp>
#include <boost/lexical_cast.hpp>
#include <boost/scoped_ptr.hpp>
#include <fcntl.h>
#include <fstream>
#include <set>

#include "mongo/base/init.h"
#include "mongo/bson/util/bson_extract.h"
#include "mongo/client/auth_helpers.h"
#include "mongo/client/dbclientcursor.h"
#include "mongo/db/auth/authorization_manager.h"
#include "mongo/db/auth/authorization_manager_global.h"
#include "mongo/db/auth/authz_manager_external_state_d.h"
#include "mongo/db/auth/user_name.h"
#include "mongo/db/auth/role_name.h"
#include "mongo/db/json.h"
#include "mongo/db/namespace_string.h"
#include "mongo/tools/mongorestore_options.h"
#include "mongo/tools/tool.h"
#include "mongo/util/mmap.h"
#include "mongo/util/options_parser/option_section.h"
#include "mongo/util/stringutils.h"

using namespace mongo;

namespace {
    const char* OPLOG_SENTINEL = "$oplog";  // compare by ptr not strcmp
}

MONGO_INITIALIZER_WITH_PREREQUISITES(RestoreAuthExternalState, ("ToolMocks"))(
        InitializerContext* context) {
    // Give restore the mongod implementation of AuthorizationManager so that it can run
    // the _mergeAuthzCollections command directly against the data files
    clearGlobalAuthorizationManager();
    setGlobalAuthorizationManager(new AuthorizationManager(
            new AuthzManagerExternalStateMongod()));

    return Status::OK();
}

class Restore : public BSONTool {
public:

    string _curns;
    string _curdb;
    string _curcoll;
    string _serverBinVersion; // Version identifier of the server we're restoring to
    set<UserName> _users; // Holds users that are already in the cluster when restoring with --drop
    set<RoleName> _roles; // Holds roles that are already in the cluster when restoring with --drop
    scoped_ptr<Matcher> _opmatcher; // For oplog replay
    scoped_ptr<OpTime> _oplogLimitTS; // for oplog replay (limit)
    int _oplogEntrySkips; // oplog entries skipped
    int _oplogEntryApplies; // oplog entries applied
    int _serverAuthzVersion; // authSchemaVersion of the cluster being restored into.
    int _dumpFileAuthzVersion; // version extracted from admin.system.version file in dump.
    bool _serverAuthzVersionDocExists; // Whether the remote cluster has an admin.system.version doc
    Restore() : BSONTool(), _oplogEntrySkips(0), _oplogEntryApplies(0), _serverAuthzVersion(0),
            _dumpFileAuthzVersion(0), _serverAuthzVersionDocExists(false) { }

    virtual void printHelp(ostream& out) {
        printMongoRestoreHelp(&out);
    }

    void storeRemoteAuthzVersion() {
        Status status = auth::getRemoteStoredAuthorizationVersion(&conn(),
                                                                  &_serverAuthzVersion);
        uassertStatusOK(status);
        uassert(17370,
                mongoutils::str::stream() << "Restoring users and roles is only supported for "
                        "clusters with auth schema versions " <<
                        AuthorizationManager::schemaVersion24 << " or " <<
                        AuthorizationManager::schemaVersion26Final << ", found: " <<
                        _serverAuthzVersion,
                _serverAuthzVersion == AuthorizationManager::schemaVersion24 ||
                _serverAuthzVersion == AuthorizationManager::schemaVersion26Final);

        _serverAuthzVersionDocExists = !conn().findOne(
                AuthorizationManager::versionCollectionNamespace,
                AuthorizationManager::versionDocumentQuery).isEmpty();
    }

    virtual int doRun() {

        boost::filesystem::path root = mongoRestoreGlobalParams.restoreDirectory;

        // check if we're actually talking to a machine that can write
        if (!isMaster()) {
            return -1;
        }

        if (isMongos() && toolGlobalParams.db == "" && exists(root / "config")) {
            toolError() << "Cannot do a full restore on a sharded system" << std::endl;
            return -1;
        }

        {
            // Store server's version
            BSONObj out;
            if (! conn().simpleCommand("admin", &out, "buildinfo")) {
                toolError() << "buildinfo command failed: "
                        << out["errmsg"].String() << std::endl;
                return -1;
            }

            _serverBinVersion = out["version"].String();
        }

        if (mongoRestoreGlobalParams.restoreUsersAndRoles) {
            storeRemoteAuthzVersion(); // populate _serverAuthzVersion

            if (_serverAuthzVersion == AuthorizationManager::schemaVersion26Final) {
                uassert(17410,
                        mongoutils::str::stream() << mongoRestoreGlobalParams.tempUsersColl <<
                                " collection already exists, but is needed to restore user data.  "
                                "Drop this collection or specify a different collection (via "
                                "--tempUsersColl) to use to temporarily hold user data during the "
                                "restore process",
                        !conn().exists(mongoRestoreGlobalParams.tempUsersColl));
                uassert(17411,
                        mongoutils::str::stream() << mongoRestoreGlobalParams.tempRolesColl <<
                                " collection already exists, but is needed to restore role data.  "
                                "Drop this collection or specify a different collection (via "
                                "--tempRolesColl) to use to temporarily hold role data during the "
                                "restore process",
                        !conn().exists(mongoRestoreGlobalParams.tempRolesColl));
            }

            if (toolGlobalParams.db.empty() && toolGlobalParams.coll.empty() &&
                    exists(root / "admin" / "system.version.bson")) {
                // Will populate _dumpFileAuthzVersion
                processFileAndMetadata(root / "admin" / "system.version.bson",
                                       "admin.system.version");
            } else if (!toolGlobalParams.db.empty()) {
                // DB-specific restore
                if (exists(root / "$admin.system.users.bson")) {
                    uassert(17372,
                            mongoutils::str::stream() << "$admin.system.users.bson file found, "
                                    "which implies that the dump was taken from a system with "
                                    "schema version " << AuthorizationManager::schemaVersion26Final
                                    << " users, but server has authorization schema version "
                                    << _serverAuthzVersion,
                            _serverAuthzVersion == AuthorizationManager::schemaVersion26Final);
                    toolInfoLog() << "Restoring users for the " << toolGlobalParams.db <<
                            " database to admin.system.users" << endl;
                    processFileAndMetadata(root / "$admin.system.users.bson", "admin.system.users");
                }
                if (exists(root / "$admin.system.roles.bson")) {
                    uassert(17373,
                            mongoutils::str::stream() << "$admin.system.roles.bson file found, "
                                    "which implies that the dump was taken from a system with  "
                                    "schema version " << AuthorizationManager::schemaVersion26Final
                                    << " authorization data, but server has authorization schema "
                                    "version " << _serverAuthzVersion,
                            _serverAuthzVersion == AuthorizationManager::schemaVersion26Final);
                    toolInfoLog() << "Restoring roles for the " << toolGlobalParams.db <<
                            " database to admin.system.roles" << endl;
                    processFileAndMetadata(root / "$admin.system.roles.bson", "admin.system.roles");
                }
            }
        }

        if (mongoRestoreGlobalParams.oplogReplay) {
            // fail early if errors

            if (toolGlobalParams.db != "") {
                toolError() << "Can only replay oplog on full restore" << std::endl;
                return -1;
            }

            if ( ! exists(root / "oplog.bson") ) {
                toolError() << "No oplog file to replay. Make sure you run mongodump with --oplog."
                          << std::endl;
                return -1;
            }

            if (versionCmp(_serverBinVersion, "1.7.4-pre-") < 0) {
                toolError() << "Can only replay oplog to server version >= 1.7.4" << std::endl;
                return -1;
            }

            string oplogInc = "0";

            if(!mongoRestoreGlobalParams.oplogLimit.empty()) {
                size_t i = mongoRestoreGlobalParams.oplogLimit.find_first_of(':');
                if ( i != string::npos ) {
                    if (i + 1 < mongoRestoreGlobalParams.oplogLimit.length()) {
                        oplogInc = mongoRestoreGlobalParams.oplogLimit.substr(i + 1);
                    }

                    mongoRestoreGlobalParams.oplogLimit =
                        mongoRestoreGlobalParams.oplogLimit.substr(0, i);
                }

                try {
                    _oplogLimitTS.reset(new OpTime(
                        boost::lexical_cast<unsigned long>(
                            mongoRestoreGlobalParams.oplogLimit.c_str()),
                        boost::lexical_cast<unsigned long>(oplogInc.c_str())));
                } catch( const boost::bad_lexical_cast& ) {
                    toolError() << "Could not parse oplogLimit into Timestamp from values ( "
                              << mongoRestoreGlobalParams.oplogLimit << " , " << oplogInc << " )"
                              << std::endl;
                    return -1;
                }

                if (!mongoRestoreGlobalParams.oplogLimit.empty()) {
                    // Only for a replica set as master will have no-op entries so we would need to
                    // skip them all to find the real op
                    scoped_ptr<DBClientCursor> cursor(
                            conn().query("local.oplog.rs", Query().sort(BSON("$natural" << -1)),
                                         1 /*return first*/));
                    OpTime tsOptime;
                    // get newest oplog entry and make sure it is older than the limit to apply.
                    if (cursor->more()) {
                        tsOptime = cursor->next().getField("ts")._opTime();
                        if (tsOptime > *_oplogLimitTS.get()) {
                            toolError() << "The oplogLimit is not newer than"
                                      << " the last oplog entry on the server."
                                      << std::endl;
                            return -1;
                        }
                    }

                    BSONObjBuilder tsRestrictBldr;
                    if (!tsOptime.isNull())
                        tsRestrictBldr << "$gt" << tsOptime;
                    tsRestrictBldr << "$lt" << *_oplogLimitTS.get();

                    BSONObj query = BSON("ts" << tsRestrictBldr.obj());

                    if (!tsOptime.isNull()) {
                        toolInfoLog() << "Latest oplog entry on the server is " << tsOptime.getSecs()
                                      << ":" << tsOptime.getInc() << std::endl;
                        toolInfoLog() << "Only applying oplog entries matching this criteria: "
                                      << query.jsonString() << std::endl;
                    }
                    _opmatcher.reset(new Matcher(query,
                                                 MatchExpressionParser::WhereCallback()));
                }
            }
        }

        /* If toolGlobalParams.db is not "" then the user specified a db name to restore as.
         *
         * In that case we better be given either a root directory that
         * contains only .bson files or a single .bson file  (a db).
         *
         * In the case where a collection name is specified we better be
         * given either a root directory that contains only a single
         * .bson file, or a single .bson file itself (a collection).
         */
        drillDown(root, toolGlobalParams.db != "", toolGlobalParams.coll != "",
                  !(_oplogLimitTS.get() == NULL), true);

        // should this happen for oplog replay as well?
        string err = conn().getLastError(toolGlobalParams.db == "" ? "admin" : toolGlobalParams.db);
        if (!err.empty()) {
            toolError() << err << std::endl;
        }

        if (mongoRestoreGlobalParams.oplogReplay) {
            toolInfoLog() << "\t Replaying oplog" << std::endl;
            _curns = OPLOG_SENTINEL;
            processFile( root / "oplog.bson" );
            toolInfoLog() << "Applied " << _oplogEntryApplies << " oplog entries out of "
                          << _oplogEntryApplies + _oplogEntrySkips << " (" << _oplogEntrySkips
                          << " skipped)." << std::endl;
        }

        return EXIT_CLEAN;
    }

    void drillDown( boost::filesystem::path root,
                    bool use_db,
                    bool use_coll,
                    bool oplogReplayLimit,
                    bool top_level=false) {
        bool json_metadata = false;
        if (logger::globalLogDomain()->shouldLog(logger::LogSeverity::Debug(2))) {
            toolInfoLog() << "drillDown: " << root.string() << std::endl;
        }

        // skip hidden files and directories
        if (root.leaf().string()[0] == '.' && root.leaf().string() != ".")
            return;

        if ( is_directory( root ) ) {
            boost::filesystem::directory_iterator end;
            boost::filesystem::directory_iterator i(root);
            boost::filesystem::path indexes;
            while ( i != end ) {
                boost::filesystem::path p = *i;
                i++;

                if (use_db) {
                    if (boost::filesystem::is_directory(p)) {
                        toolError() << "ERROR: root directory must be a dump of a single database"
                                  << std::endl;
                        toolError() << "       when specifying a db name with --db" << std::endl;
                        toolError() << "       use the --help option for more information"
                                  << std::endl;
                        return;
                    }
                }

                if (use_coll) {
                    if (boost::filesystem::is_directory(p) || i != end) {
                        toolError() << "ERROR: root directory must be a dump of a single collection"
                                  << std::endl;
                        toolError() << "       when specifying a collection name with --collection"
                                  << std::endl;
                        toolError() << "       use the --help option for more information"
                                  << std::endl;
                        return;
                    }
                }

                // Ignore system.indexes.bson if we have *.metadata.json files
                if ( endsWith( p.string().c_str() , ".metadata.json" ) ) {
                    json_metadata = true;
                }

                // don't insert oplog
                if (top_level && !use_db && p.leaf() == "oplog.bson")
                    continue;

                if ( p.leaf() == "system.indexes.bson" ) {
                    indexes = p;
                } else {
                    drillDown(p, use_db, use_coll, oplogReplayLimit);
                }
            }

            if (!indexes.empty() && !json_metadata) {
                drillDown(indexes, use_db, use_coll, oplogReplayLimit);
            }

            return;
        }

        if (oplogReplayLimit) {
            toolError() << "The oplogLimit option cannot be used if "
                      << "normal databases/collections exist in the dump directory."
                      << std::endl;
            exit(EXIT_FAILURE);
        }

        string ns;
        if (use_db) {
            ns += toolGlobalParams.db;
        }
        else {
            ns = root.parent_path().filename().string();
            if (ns.empty())
                ns = "test";
        }

        verify( ns.size() );

        string oldCollName = root.leaf().string(); // Name of the collection that was dumped from
        oldCollName = oldCollName.substr( 0 , oldCollName.find_last_of( "." ) );
        if (use_coll) {
            ns += "." + toolGlobalParams.coll;
        }
        else {
            ns += "." + oldCollName;
        }

        if ( endsWith( root.string().c_str() , ".metadata.json" ) ) {
            // Metadata files are handled when the corresponding .bson file is handled
            return;
        }

        if ((root.leaf() == "system.version.bson" && toolGlobalParams.db.empty()) ||
                root.leaf() == "$admin.system.users.bson" ||
                root.leaf() == "$admin.system.roles.bson") {
            // These files were already explicitly handled at the beginning of the restore.
            return;
        }

        boost::filesystem::file_status fileStatus = boost::filesystem::status(root);
        if ( ! ( endsWith( root.string().c_str() , ".bson" ) ||
                 endsWith( root.string().c_str() , ".bin" ) ) &&
             ! ( root.string() == "-" ) &&
             ! ( fileStatus.type() == boost::filesystem::fifo_file ) ) {
            toolError() << "don't know what to do with file [" << root.string() << "]" << std::endl;
            return;
        }

        // Both --db and --collection have to be provided when using stdin or fifo.
        if ((root.string() == "-" || fileStatus.type() == boost::filesystem::fifo_file) &&
            !(use_db && use_coll)) {
            toolError() << "Both --db and --collection have to be provided when using stding/fifo";
            exit(EXIT_FAILURE);
        }

        toolInfoLog() << root.string() << std::endl;

        if ( root.leaf() == "system.profile.bson" ) {
            toolInfoLog() << "\t skipping system.profile.bson" << std::endl;
            return;
        }

        processFileAndMetadata(root, ns);
    }

    /**
     * 1) Drop collection if --drop was specified.  For system.users or system.roles collections,
     * however, you don't want to remove all the users/roles up front as some of them may be needed
     * by the restore.  Instead, keep a set of all the users/roles originally in the server, then
     * after restoring the users/roles from the dump, remove any users roles that were present in
     * the system originally but aren't in the dump.
     *
     * 2) Parse metadata file (if present) and if the collection doesn't exist (or was just dropped
     * b/c we're using --drop), create the collection with the options from the metadata file
     *
     * 3) Restore the data from the dump file for this collection
     *
     * 4) If the user asked to drop this collection, then at this point the _users and _roles sets
     * will contain users and roles that were in the collection but not in the dump we are
     * restoring. Iterate these sets and delete any users and roles that are there.
     *
     * 5) Restore indexes based on index definitions from the metadata file.
     */
    void processFileAndMetadata(const boost::filesystem::path& root, const std::string& ns) {

        _curns = ns;
        _curdb = nsToDatabase(_curns);
        _curcoll = nsToCollectionSubstring(_curns).toString();

        toolInfoLog() << "\tgoing into namespace [" << _curns << "]" << std::endl;

        // 1) Drop collection if needed.  Save user and role data if this is a system.users or
        // system.roles collection
        if (mongoRestoreGlobalParams.drop) {
            if (_curcoll == "system.users") {
                if (_serverAuthzVersion == AuthorizationManager::schemaVersion24 ||
                            _curdb != "admin") {
                    // Restoring 2.4-style user docs so can't use the _mergeAuthzCollections command
                    // Create map of the users currently in the DB so the ones that don't show up in
                    // the dump file can be removed later.
                    BSONObj fields = BSON("user" << 1 << "userSource" << 1);
                    scoped_ptr<DBClientCursor> cursor(conn().query(_curns, Query(), 0, 0, &fields));
                    while (cursor->more()) {
                        BSONObj user = cursor->next();
                        string userDB;
                        uassertStatusOK(bsonExtractStringFieldWithDefault(user,
                                                                          "userSource",
                                                                          _curdb,
                                                                          &userDB));
                        _users.insert(UserName(user["user"].String(), userDB));
                    }
                }
            }
            else if (!startsWith(_curcoll, "system.")) { // Can't drop system collections
                toolInfoLog() << "\t dropping" << std::endl;
                conn().dropCollection( ns );
            }
        } else {
            // If drop is not used, warn if the collection exists.
            scoped_ptr<DBClientCursor> cursor(conn().query(_curdb + ".system.namespaces",
                                                           Query(BSON("name" << ns))));
            if (cursor->more()) {
                // collection already exists show warning
                toolError() << "Restoring to " << ns << " without dropping. Restored data "
                        << "will be inserted without raising errors; check your server log"
                        << std::endl;
            }
        }

        // 2) Create collection with options from metadata file if present
        BSONObj metadataObject;
        boost::filesystem::file_status fileStatus = boost::filesystem::status(root);
        if (fileStatus.type() != boost::filesystem::fifo_file &&
            root.string() != "-" &&
            (mongoRestoreGlobalParams.restoreOptions || mongoRestoreGlobalParams.restoreIndexes)) {
            string oldCollName = root.leaf().string(); // Name of collection that was dumped from
            oldCollName = oldCollName.substr( 0 , oldCollName.find_last_of( "." ) );
            boost::filesystem::path metadataFile = (root.branch_path() / (oldCollName + ".metadata.json"));
            if (!boost::filesystem::exists(metadataFile.string())) {
                // This is fine because dumps from before 2.1 won't have a metadata file, just print a warning.
                // System collections shouldn't have metadata so don't warn if that file is missing.
                if (!startsWith(metadataFile.leaf().string(), "system.")) {
                    toolInfoLog() << metadataFile.string() << " not found. Skipping." << std::endl;
                }
            } else {
                metadataObject = parseMetadataFile(metadataFile.string());
            }
        }

        if (mongoRestoreGlobalParams.restoreOptions && metadataObject.hasField("options")) {
            // Try to create collection with given options
            createCollectionWithOptions(metadataObject["options"].Obj());
        }

        // 3) Actually restore the BSONObjs inside the dump file
        processFile( root );

        // 4) If running with --drop, remove any users/roles that were in the system at the
        // beginning of the restore but weren't found in the dump file
        if (_curcoll == "system.users") {
            if ((_serverAuthzVersion == AuthorizationManager::schemaVersion24 ||
                    _curdb != "admin")) {
                // Restoring 2.4 style user docs so don't use the _mergeAuthzCollections command
                if (mongoRestoreGlobalParams.drop) {
                    // Delete any users that used to exist but weren't in the dump file
                    for (set<UserName>::iterator it = _users.begin(); it != _users.end(); ++it) {
                        const UserName& name = *it;
                        BSONObjBuilder queryBuilder;
                        queryBuilder << "user" << name.getUser();
                        if (name.getDB() == _curdb) {
                            // userSource field won't be present for v1 users docs in the same db as
                            // the user is defined on.
                            queryBuilder << "userSource" << BSONNULL;
                        } else {
                            queryBuilder << "userSource" << name.getDB();
                        }
                        conn().remove(_curns, Query(queryBuilder.done()));
                    }
                    _users.clear();
                }
            } else {
                // Use _mergeAuthzCollections command to move into admin.system.users the user
                // docs that were restored into the temp user collection
                BSONObj res;
                conn().runCommand("admin",
                                  BSON("_mergeAuthzCollections" << 1 <<
                                       "tempUsersCollection" <<
                                               mongoRestoreGlobalParams.tempUsersColl <<
                                       "drop" << mongoRestoreGlobalParams.drop <<
                                       "writeConcern" << BSON("w" << mongoRestoreGlobalParams.w)),
                                  res);
                uassert(17412,
                        mongoutils::str::stream() << "Cannot restore users because the "
                                "_mergeAuthzCollections command failed: " << res.toString(),
                        res["ok"].trueValue());

                conn().dropCollection(mongoRestoreGlobalParams.tempUsersColl);
            }
        }
        if (_curns == "admin.system.roles") {
            // Use _mergeAuthzCollections command to move into admin.system.roles the role
            // docs that were restored into the temp roles collection
            BSONObj res;
            conn().runCommand("admin",
                              BSON("_mergeAuthzCollections" << 1 <<
                                   "tempRolesCollection" <<
                                           mongoRestoreGlobalParams.tempRolesColl <<
                                   "drop" << mongoRestoreGlobalParams.drop <<
                                   "writeConcern" << BSON("w" << mongoRestoreGlobalParams.w)),
                              res);
            uassert(17413,
                    mongoutils::str::stream() << "Cannot restore roles because the "
                            "_mergeAuthzCollections command failed: " << res.toString(),
                    res["ok"].trueValue());

            conn().dropCollection(mongoRestoreGlobalParams.tempRolesColl);
        }

        // 5) Restore indexes
        if (mongoRestoreGlobalParams.restoreIndexes && metadataObject.hasField("indexes")) {
            vector<BSONElement> indexes = metadataObject["indexes"].Array();
            for (vector<BSONElement>::iterator it = indexes.begin(); it != indexes.end(); ++it) {
                createIndex((*it).Obj(), false);
            }
        }
    }

    virtual void gotObject( const BSONObj& obj ) {
        if (_curns == OPLOG_SENTINEL) { // intentional ptr compare
            if (obj["op"].valuestr()[0] == 'n') // skip no-ops
                return;

            // exclude operations that don't meet (timestamp) criteria
            if ( _opmatcher.get() && ! _opmatcher->matches ( obj ) ) {
                _oplogEntrySkips++;
                return;
            }

            string db = obj["ns"].valuestr();
            db = db.substr(0, db.find('.'));

            BSONObj cmd = BSON( "applyOps" << BSON_ARRAY( obj ) );
            BSONObj out;
            conn().runCommand(db, cmd, out);
            _oplogEntryApplies++;

            // wait for ops to propagate to "w" nodes (doesn't warn if w used without replset)
            if (mongoRestoreGlobalParams.w > 0) {
                string err = conn().getLastError(db, false, false, mongoRestoreGlobalParams.w);
                if (!err.empty()) {
                    toolError() << "Error while replaying oplog: " << err << std::endl;
                }
            }
            return;
        }

        if (nsToCollectionSubstring(_curns) == "system.indexes") {
            createIndex(obj, true);
        }
        else if (_curns == "admin.system.roles") {
            // To prevent modifying roles when other role modifications may be going on, restore
            // the roles to a temporary collection and merge them into admin.system.roles later
            // using the _mergeAuthzCollections command.
            conn().insert(mongoRestoreGlobalParams.tempRolesColl, obj);
        }
        else if (_curcoll == "system.users") {
            uassert(17416,
                    mongoutils::str::stream() << "Cannot modify user data on a server with version "
                            "greater than or equal to 2.5.4 that has not yet updated the "
                            "authorization data to schema version " <<
                            AuthorizationManager::schemaVersion26Final <<
                            ". Found server version " << _serverBinVersion << " with "
                            "authorization schema version " << _serverAuthzVersion,
                    _curdb != "admin" ||
                    versionCmp(_serverBinVersion, "2.5.4") < 0 ||
                    _serverAuthzVersion == AuthorizationManager::schemaVersion26Final);

            if (obj.hasField("credentials")) {
                if (_serverAuthzVersion == AuthorizationManager::schemaVersion24) {
                    // v3 user, v1 system
                    uasserted(17407,
                              mongoutils::str::stream()
                                      << "Server has authorization schema version "
                                      << AuthorizationManager::schemaVersion24
                                      << ", but found a schema version "
                                      << AuthorizationManager::schemaVersion26Final << " user: "
                                      << obj.toString());
                } else {
                    // v3 user, v3 system
                    uassert(17414,
                            mongoutils::str::stream() << "Found a schema version " <<
                                    AuthorizationManager::schemaVersion26Final <<
                                    " user when restoring to a non-admin db system.users "
                                    "collection: " << obj.toString(),
                            _curdb == "admin");
                    // To prevent modifying users when other user modifications may be going on,
                    // restore the users to a temporary collection and merge them into
                    // admin.system.users later using the _mergeAuthzCollections command.
                    conn().insert(mongoRestoreGlobalParams.tempUsersColl, obj);
                }
            } else {
                if (_curdb == "admin" &&
                        _serverAuthzVersion == AuthorizationManager::schemaVersion26Final &&
                        !_serverAuthzVersionDocExists) {
                    // server with schemaVersion26Final implies it is running 2.5.4 or greater.
                    uasserted(17415,
                              mongoutils::str::stream() << "Cannot restore users with schema " <<
                                      "version " << AuthorizationManager::schemaVersion24 <<
                                      " to a system with server version 2.5.4 or greater");
                }

                if (_serverAuthzVersion == AuthorizationManager::schemaVersion24 ||
                        _curdb != "admin") { // Restoring 2.4 schema users to non-admin dbs is OK)
                    // v1 user, v1 system
                    string userDB;
                    uassertStatusOK(bsonExtractStringFieldWithDefault(obj,
                                                                      "userSource",
                                                                      _curdb,
                                                                      &userDB));

                    if (mongoRestoreGlobalParams.drop && _users.count(UserName(obj["user"].String(),
                                                                               userDB))) {
                        // Since system collections can't be dropped, we have to manually
                        // replace the contents of the system.users collection
                        BSONObj userMatch = BSON("user" << obj["user"].String() <<
                                                 "userSource" << userDB);
                        conn().update(_curns, Query(userMatch), obj);
                        _users.erase(UserName(obj["user"].String(), userDB));
                    } else {
                        conn().insert(_curns, obj);
                    }
                } else {
                    // v1 user, v3 system
                    // TODO(spencer): SERVER-12491 Rather than failing here, we should convert the
                    // v1 user to an equivalent v3 schema user
                    uasserted(17408,
                              mongoutils::str::stream()
                                      << "Server has authorization schema version "
                                      << AuthorizationManager::schemaVersion26Final
                                      << ", but found a schema version "
                                      << AuthorizationManager::schemaVersion24 << " user: "
                                      << obj.toString());
                }
            }
        }
        else {
            if (_curns == "admin.system.version") {
                long long authVersion;
                uassertStatusOK(bsonExtractIntegerField(obj,
                                                        AuthorizationManager::schemaVersionFieldName,
                                                        &authVersion));
                _dumpFileAuthzVersion = static_cast<int>(authVersion);
                uassert(17371,
                        mongoutils::str::stream() << "Server's authorization data schema version "
                                "does not match that of the data in the dump file.  Server's schema"
                                " version: " << _serverAuthzVersion << ", schema version in dump: "
                                << _dumpFileAuthzVersion,
                    _serverAuthzVersion == _dumpFileAuthzVersion);
            }
            conn().insert(_curns, obj);
        }

        // wait for insert (or update) to propagate to "w" nodes (doesn't warn if w used
        // without replset)
        if (mongoRestoreGlobalParams.w > 0) {
            string err = conn().getLastError(_curdb, false, false, mongoRestoreGlobalParams.w);
            if (!err.empty()) {
                toolError() << err << std::endl;
            }
        }
    }

private:

    BSONObj parseMetadataFile(string filePath) {
        long long fileSize = boost::filesystem::file_size(filePath);
        ifstream file(filePath.c_str(), ios_base::in);

        boost::scoped_array<char> buf(new char[fileSize + 1]);
        file.read(buf.get(), fileSize);
        buf[fileSize] = '\0';

        int objSize;
        BSONObj obj;
        obj = fromjson(buf.get(), &objSize);
        return obj;
    }

    // Compares 2 BSONObj representing collection options. Returns true if the objects
    // represent different options. Ignores the "create" field.
    bool optionsSame(BSONObj obj1, BSONObj obj2) {
        int nfields = 0;
        BSONObjIterator i(obj1);
        while ( i.more() ) {
            BSONElement e = i.next();
            if (!obj2.hasField(e.fieldName())) {
                if (strcmp(e.fieldName(), "create") == 0) {
                    continue;
                } else {
                    return false;
                }
            }
            nfields++;
            if (e != obj2[e.fieldName()]) {
                return false;
            }
        }
        return nfields == obj2.nFields();
    }

    void createCollectionWithOptions(BSONObj obj) {
        BSONObjIterator i(obj);

        // Rebuild obj as a command object for the "create" command.
        // - {create: <name>} comes first, where <name> is the new name for the collection
        // - elements with type Undefined get skipped over
        BSONObjBuilder bo;
        bo.append("create", _curcoll);
        while (i.more()) {
            BSONElement e = i.next();

            if (strcmp(e.fieldName(), "create") == 0) {
                continue;
            }

            if (e.type() == Undefined) {
                toolInfoLog() << _curns << ": skipping undefined field: " << e.fieldName()
                              << std::endl;
                continue;
            }

            bo.append(e);
        }
        obj = bo.obj();

        BSONObj fields = BSON("options" << 1);
        scoped_ptr<DBClientCursor> cursor(conn().query(_curdb + ".system.namespaces", Query(BSON("name" << _curns)), 0, 0, &fields));

        bool createColl = true;
        if (cursor->more()) {
            createColl = false;
            BSONObj nsObj = cursor->next();
            if (!nsObj.hasField("options") || !optionsSame(obj, nsObj["options"].Obj())) {
                toolError() << "WARNING: collection " << _curns
                          << " exists with different options than are in the metadata.json file and"
                          << " not using --drop. Options in the metadata file will be ignored."
                          << std::endl;
            }
        }

        if (!createColl) {
            return;
        }

        BSONObj info;
        if (!conn().runCommand(_curdb, obj, info)) {
            uasserted(15936, "Creating collection " + _curns + " failed. Errmsg: " + info["errmsg"].String());
        } else {
            toolInfoLog() << "\tCreated collection " << _curns << " with options: "
                          << obj.jsonString() << std::endl;
        }
    }

    /* We must handle if the dbname or collection name is different at restore time than what was dumped.
       If keepCollName is true, however, we keep the same collection name that's in the index object.
     */
    void createIndex(BSONObj indexObj, bool keepCollName) {
        BSONObjBuilder bo;
        BSONObjIterator i(indexObj);
        while ( i.more() ) {
            BSONElement e = i.next();
            if (strcmp(e.fieldName(), "ns") == 0) {
                NamespaceString n(e.String());
                string s = _curdb + "." + (keepCollName ? n.coll().toString() : _curcoll);
                bo.append("ns", s);
            }
            // Remove index version number
            else if (strcmp(e.fieldName(), "v") != 0 || mongoRestoreGlobalParams.keepIndexVersion) {
                bo.append(e);
            }
        }
        BSONObj o = bo.obj();
        if (logger::globalLogDomain()->shouldLog(logger::LogSeverity::Debug(0))) {
            toolInfoLog() << "\tCreating index: " << o << std::endl;
        }
        conn().insert( _curdb + ".system.indexes" ,  o );

        // We're stricter about errors for indexes than for regular data
        BSONObj err = conn().getLastErrorDetailed(_curdb, false, false, mongoRestoreGlobalParams.w);

        if (err.hasField("err") && !err["err"].isNull()) {
            if (err["err"].str() == "norepl" && mongoRestoreGlobalParams.w > 1) {
                toolError() << "Cannot specify write concern for non-replicas" << std::endl;
            }
            else {
                string errCode;

                if (err.hasField("code")) {
                    errCode = str::stream() << err["code"].numberInt();
                }

                toolError() << "Error creating index " << o["ns"].String() << ": "
                          << errCode << " " << err["err"] << std::endl;
            }

            ::abort();
        }

        massert(16441, str::stream() << "Error calling getLastError: " << err["errmsg"],
                err["ok"].trueValue());
    }
};

REGISTER_MONGO_TOOL(Restore);

max_stars_repo_name: EshaMaharishi/mongo

max_issues_repo_name: EshaMaharishi/mongo

max_forks_repo_name: EshaMaharishi/mongo


text:
//          Copyright Gavin Band 2008 - 2012.
// Distributed under the Boost Software License, Version 1.0.
//    (See accompanying file LICENSE_1_0.txt or copy at
//          http://www.boost.org/LICENSE_1_0.txt)

#include <iostream>
#include <iomanip>
#include <vector>
#include <Eigen/Core>
#include "genfile/SNPDataSourceProcessor.hpp"
#include "genfile/string_utils.hpp"
#include "genfile/Error.hpp"
#include "genfile/vcf/get_set_eigen.hpp"
#include "components/RelatednessComponent/PCAProjector.hpp"
#include "components/RelatednessComponent/PCAComputer.hpp"
#include "components/RelatednessComponent/mean_centre_genotypes.hpp"
#include "components/RelatednessComponent/names.hpp"
#include "components/RelatednessComponent/write_matrix_to_stream.hpp"

// #define DEBUG_PCA_PROJECTOR 1

namespace pca {
	PCAProjector::UniquePtr PCAProjector::create(
		genfile::CohortIndividualSource const& samples,
		appcontext::UIContext& ui_context,
		genfile::VariantIdentifyingData::CompareFields const& comparer
	) {
		return PCAProjector::UniquePtr( new PCAProjector( samples, ui_context, comparer )) ;
	}

	PCAProjector::PCAProjector(
		genfile::CohortIndividualSource const& samples,
		appcontext::UIContext& ui_context,
		genfile::VariantIdentifyingData::CompareFields const& comparer
	):
		m_samples( samples ),
	 	m_ui_context( ui_context ),
		m_snps( comparer )
	{}

	void PCAProjector::set_loadings(
		std::vector< genfile::VariantIdentifyingData > const& snps,
		Vector const& counts,
		Vector const& frequencies,
		Matrix const& loadings,
		std::vector< std::string > const& names
	) {
		using genfile::string_utils::to_string ;
		if( snps.size() != std::size_t( loadings.rows() ) ) {
			throw genfile::MismatchError(
				"relatedness::PCAProjector::set_loadings()",
				"(function arguments)",
				"snp.size()=" + to_string( snps.size() ),
				"loadings.rows()=" + to_string( loadings.rows() )
			) ;
		}
		if( names.size() != std::size_t( loadings.cols() ) ) {
			throw genfile::MismatchError(
				"relatedness::PCAProjector::set_loadings()",
				"(function arguments)",
				"names.size()=" + to_string( names.size() ),
				"loadings.cols()=" + to_string( loadings.cols() )
			) ;
		}
		for( std::size_t i = 0; i < snps.size(); ++i ) {
			std::pair< SnpMap::const_iterator, bool > where = m_snps.insert( std::make_pair( snps[i], int( i ) ) ) ;
			if( !where.second ) {
				throw genfile::DuplicateKeyError( "snps argument to pca::PCAProjector::set_loadings()", to_string( snps[i] )) ;
			}
		}
		m_counts = counts ;
		m_frequencies = frequencies ;
		m_loadings = loadings ;
		m_names = names ;

#if DEBUG_PCA_PROJECTOR
		std::cerr << "PCAProjector: loadings are:\n" ;
		pca::write_matrix_to_stream( std::cerr, m_loadings.block( 0, 0, 10, m_loadings.cols() )) ;
#endif
	}

	void PCAProjector::send_results_to( ResultSignal::slot_type callback ) {
		m_result_signal.connect( callback ) ;
	}

	void PCAProjector::begin_processing_snps( std::size_t number_of_samples, genfile::SNPDataSource::Metadata const& ) {
		m_projections.setZero( number_of_samples, m_loadings.cols() ) ;
		m_snps_visited_per_sample = Eigen::VectorXd::Zero( number_of_samples ) ;
		m_total_snps_visited = 0 ;
	}

	void PCAProjector::processed_snp( genfile::VariantIdentifyingData const& snp, genfile::VariantDataReader& data_reader ) {
		SnpMap::const_iterator where = m_snps.find( snp ) ;
		if( where != m_snps.end() && ( m_loadings.row( where->second ).sum() == m_loadings.row( where->second ).sum() ) ) {
			data_reader.get(
				":genotypes:",
				genfile::vcf::get_threshholded_calls( m_genotype_calls, m_non_missingness, 0.9, 0, 0, 1, 2 )
			) ;
			assert( m_genotype_calls.size() == m_projections.rows() ) ;
			assert( m_non_missingness.size() == m_genotype_calls.size() ) ;
			double const allele_frequency = m_frequencies( where->second ) ;
			double const maf = std::min( allele_frequency, 1 - allele_frequency ) ;
			double const count = m_counts( where->second ) ;
			if( count > 0 && maf > 0.001 ) {
#if DEBUG_PCA_PROJECTOR
				std::cerr << std::resetiosflags( std::ios::floatfield ) << std::setprecision( 5 ) ;
				std::cerr << "SNP: " << snp << ": freq = " << allele_frequency << ", uncentred genotypes are: " << m_genotype_calls.transpose().head( 20 ) << "...\n" ;
#endif
				pca::mean_centre_genotypes( &m_genotype_calls, m_non_missingness, allele_frequency ) ;
				m_genotype_calls /= std::sqrt( 2.0 * allele_frequency * ( 1.0 - allele_frequency ) ) ;
#if DEBUG_PCA_PROJECTOR
				std::cerr << "mean-centred genotypes are: " << m_genotype_calls.transpose().head( 20 ) << "...\n" ;
				std::cerr << "non-missingness is: " << m_non_missingness.transpose().head( 20 ) << "...\n" ;
#endif

				//
				// We have loadings for C components as columns of m_loadings.
				// We have calls for N samples as entries of m_genotype_calls
				// We wish to get a matrix with samples as rows and projections as columns.
				// We therefore compute a Kronecker product of the column vectors of genotype calls
				// with the row vector of loadings for that SNP.
				//
				m_projections += m_genotype_calls * m_loadings.row( where->second ) ;
				m_snps_visited_per_sample += m_non_missingness ;
				++m_visited[ snp.get_position() ] ;
				++m_total_snps_visited ;

				// sanity check we are not getting NaNs...
				// (It's more helpful to abort because at least that shows you where the
				// error is.  A file of NAs is not very helpful in this regard).
				assert( m_projections.sum() == m_projections.sum() ) ;
			}
		}
	}

	void PCAProjector::end_processing_snps() {
		diagnose_projection() ;
		using genfile::string_utils::to_string ;

		for( int i = 0; i < m_projections.cols(); ++i ) {
			m_projections.col(i).array() /= m_snps_visited_per_sample.array().sqrt() ;
		}

		std::cerr << m_snps_visited_per_sample.transpose() << ".\n" ;

		m_result_signal(
			"Number of SNPs in loadings: " + to_string( m_snps.size() ) + "\n" +
			"...of which " + to_string( m_total_snps_visited ) + " visited in this analysis.\n" +
			"Number of samples: " + to_string( m_projections.size() ) + "\n" +
			"Note: these projections are 1/√L times the projection of samples onto the given loadings.\n",
			m_projections.array(),
			boost::bind(
				&pca::get_concatenated_sample_ids,
				&m_samples,
				_1
			),
			boost::bind(
				&PCAProjector::get_projection_name,
				_1
			)
		) ;
	}

	genfile::VariantEntry PCAProjector::get_projection_name( std::size_t i ) {
		return std::string( "projection_onto_PC_" ) + genfile::string_utils::to_string( i + 1 ) ;
	}

	void PCAProjector::diagnose_projection() const {
		std::size_t missing_count = 0 ;
		for( SnpMap::const_iterator i = m_snps.begin(); i != m_snps.end(); ++i ) {
			VisitedSnpMap::const_iterator where = m_visited.find( i->first.get_position() ) ;
			if( where == m_visited.end() ) {
				if( missing_count == 0 ) {
					m_ui_context.logger() << "!! ( pca::PCAProjector::diagnose_projection() ): "
						<< "The following SNPs with loadings were not used in the projection computation:\n" ;
				}
				m_ui_context.logger() << "    " << i->first << "\n" ;
				++missing_count ;
			}
		}
		m_ui_context.logger() << "++ ( pca::PCAProjector::diagnose_projection() ): a total of " << m_snps.size() - missing_count
			<< " of " << m_snps.size() << " loading SNPs were used in the projection computation.\n" ;
		m_ui_context.logger() << "++ ( pca::PCAProjector::diagnose_projection() ): but wait! Total SNPs visited is " << m_total_snps_visited << ".\n" ;

		VisitedSnpMap::const_iterator i = m_visited.begin(), end_i = m_visited.end() ;
		m_ui_context.logger() << "++ ( pca::PCAProjector::diagnose_projection() ): the following positions were visited more than once:\n" ;
		bool twice = false ;
		for( ; i != end_i; ++i ) {
			if( i->second > 1 ) {
				twice = true ;
				m_ui_context.logger() << i->first << " (visited " << i->second << " time(s)).\n" ;
			}
		}
		if( !twice ) {
			m_ui_context.logger() << "(none)\n" ;
		}
	}

	std::string PCAProjector::get_metadata() const {
		using namespace genfile::string_utils ;
		return "Number of SNPs: " + to_string( m_snps.size() ) ;
	}

}

max_stars_repo_name: CreRecombinase/qctool

max_issues_repo_name: CreRecombinase/qctool

max_forks_repo_name: gavinband/qctool