/*++
Copyright (c) 2021 Microsoft Corporation

Module Name:

    polysat constraints for bit operations.

Author:

    Jakob Rath, Nikolaj Bjorner (nbjorner) 2021-12-09

Notes:

Additional possible functionality on constraints:

- activate - when operation is first activated. It may be created and only activated later.
- bit-wise assignments - narrow based on bit assignment, not entire word assignment.
- integration with congruence tables
- integration with conflict resolution

--*/

#include "math/polysat/op_constraint.h"
#include "math/polysat/solver.h"

namespace polysat {

    op_constraint::op_constraint(constraint_manager& m, code c, pdd const& p, pdd const& q, pdd const& r) :
        constraint(m, ckind_t::op_t), m_op(c), m_p(p), m_q(q), m_r(r) {
        m_vars.append(p.free_vars());
        for (auto v : q.free_vars())
            if (!m_vars.contains(v))
                m_vars.push_back(v);
        for (auto v : r.free_vars())
            if (!m_vars.contains(v))
                m_vars.push_back(v);

        switch (c) {
        case code::and_op:
            if (p.index() > q.index())
                std::swap(m_p, m_q);
            break;
        default:
            break;
        }
    }

    lbool op_constraint::eval() const {
        return eval(p(), q(), r());
    }

    lbool op_constraint::eval(assignment const& a) const {
        return eval(a.apply_to(p()), a.apply_to(q()), a.apply_to(r()));
    }

    lbool op_constraint::eval(pdd const& p, pdd const& q, pdd const& r) const {
        switch (m_op) {
        case code::lshr_op:
            return eval_lshr(p, q, r);
        case code::shl_op:
            return eval_shl(p, q, r);
        case code::and_op:
            return eval_and(p, q, r);
        default:
            return l_undef;
        }
    }

    std::ostream& op_constraint::display(std::ostream& out, lbool status) const {
        switch (status) {
        case l_true: return display(out, "==");
        case l_false: return display(out, "!=");
        default: return display(out, "?=");
        }
    }

    std::ostream& operator<<(std::ostream& out, op_constraint::code c) {
        switch (c) {
        case op_constraint::code::ashr_op:
            return out << ">>a";
        case op_constraint::code::lshr_op:
            return out << ">>";
        case op_constraint::code::shl_op:
            return out << "<<";
        case op_constraint::code::and_op:
            return out << "&";
        default:
            UNREACHABLE();
            return out;
        }
        return out;
    }

    std::ostream& op_constraint::display(std::ostream& out) const {
        return display(out, l_true);
    }

    std::ostream& op_constraint::display(std::ostream& out, char const* eq) const {
        return out << r() << " " << eq << " " << p() << " " << m_op << " " << q();
    }

    /**
     * Propagate consequences or detect conflicts based on partial assignments.
     *
     * We can assume that op_constraint is only asserted positive.
     */
    void op_constraint::narrow(solver& s, bool is_positive, bool first) {
        SASSERT(is_positive);

        if (is_currently_true(s, is_positive))
            return;

        if (first)
            activate(s);

        if (clause_ref lemma = produce_lemma(s, s.assignment()))
            s.add_clause(*lemma);
        
        if (!s.is_conflict() && is_currently_false(s, is_positive))
            s.set_conflict(signed_constraint(this, is_positive));
    }

    /**
     * Produce lemmas that contradict the given assignment.
     *
     * We can assume that op_constraint is only asserted positive.
     */
    clause_ref op_constraint::produce_lemma(solver& s, assignment const& a, bool is_positive) {
        SASSERT(is_positive);

        if (is_currently_true(a, is_positive))
            return {};

        return produce_lemma(s, a);
    }

    clause_ref op_constraint::produce_lemma(solver& s, assignment const& a) {
        switch (m_op) {
        case code::lshr_op:
            return lemma_lshr(s, a);
        case code::shl_op:
            return lemma_shl(s, a);
        case code::and_op:
            return lemma_and(s, a);
        default:
            NOT_IMPLEMENTED_YET();
            return {};
        }
    }

    void op_constraint::activate(solver& s) {
        switch (m_op) {
        case code::lshr_op:
            break;
        case code::shl_op:
            // TODO: if shift amount is constant p << k, then add p << k == p*2^k
            break;
        case code::and_op:
            // handle masking of high order bits
            activate_and(s);
            break;
        default:
            break;
        }        
    }

    unsigned op_constraint::hash() const {
        return mk_mix(p().hash(), q().hash(), r().hash());
    }

    bool op_constraint::operator==(constraint const& other) const {
        if (other.kind() != ckind_t::op_t)
            return false;
        auto const& o = other.to_op();
        return m_op == o.m_op && p() == o.p() && q() == o.q() && r() == o.r();
    }

    /**
     * Enforce basic axioms for r == p >> q:
     *
     *      q >= K  ->  r = 0
     *      q >= k  ->  r[i] = 0 for K - k <= i < K     (bit indices range from 0 to K-1, inclusive)
     *      q >= k  ->  r <= 2^{K-k} - 1
     *      q = k   ->  r[i] = p[i+k] for 0 <= i < K - k
     *      r <= p
     *      q != 0  ->  r <= p      (subsumed by previous axiom)
     *      q != 0  /\  p > 0 ->  r < p
     *      q = 0   ->  r = p
     *
     * when q is a constant, several axioms can be enforced at activation time.
     *
     * Enforce also inferences and bounds
     *
     * TODO: use also
     * s.m_viable.min_viable();
     * s.m_viable.max_viable()
     * when r, q are variables.
     */
    clause_ref op_constraint::lemma_lshr(solver& s, assignment const& a) {
        auto const pv = a.apply_to(p());
        auto const qv = a.apply_to(q());
        auto const rv = a.apply_to(r());
        unsigned const K = p().manager().power_of_2();

        signed_constraint const lshr(this, true);

        if (pv.is_val() && rv.is_val() && rv.val() > pv.val())
            // r <= p
            return s.mk_clause(~lshr, s.ule(r(), p()), true);
        else if (qv.is_val() && qv.val() >= K && rv.is_val() && !rv.is_zero())
            // q >= K -> r = 0
            return s.mk_clause(~lshr, ~s.ule(K, q()), s.eq(r()), true);
        else if (qv.is_zero() && pv.is_val() && rv.is_val() && pv != rv)
            // q = 0 -> p = r
            return s.mk_clause(~lshr, ~s.eq(q()), s.eq(p(), r()), true);
        else if (qv.is_val() && !qv.is_zero() && pv.is_val() && rv.is_val() && !pv.is_zero() && rv.val() >= pv.val())
            // q != 0 & p > 0 -> r < p
            return s.mk_clause(~lshr, s.eq(q()), s.ule(p(), 0), s.ult(r(), p()), true);
        else if (qv.is_val() && !qv.is_zero() && qv.val() < K && rv.is_val() &&
            rv.val() > rational::power_of_two(K - qv.val().get_unsigned()) - 1)
            // q >= k -> r <= 2^{K-k} - 1
            return s.mk_clause(~lshr, ~s.ule(qv.val(), q()), s.ule(r(), rational::power_of_two(K - qv.val().get_unsigned()) - 1), true);
        else if (pv.is_val() && rv.is_val() && qv.is_val() && !qv.is_zero()) {
            unsigned const k = qv.val().get_unsigned();
            // q = k  ->  r[i] = p[i+k] for 0 <= i < K - k
            for (unsigned i = 0; i < K - k; ++i) {
                if (rv.val().get_bit(i) && !pv.val().get_bit(i + k)) {
                    return s.mk_clause(~lshr, ~s.eq(q(), k), ~s.bit(r(), i), s.bit(p(), i + k), true);
                }
                if (!rv.val().get_bit(i) && pv.val().get_bit(i + k)) {
                    return s.mk_clause(~lshr, ~s.eq(q(), k), s.bit(r(), i), ~s.bit(p(), i + k), true);
                }
            }
        }
        else {
            SASSERT(!(pv.is_val() && qv.is_val() && rv.is_val()));
        }
        return {};
    }

    /** Evaluate constraint: r == p >> q */
    lbool op_constraint::eval_lshr(pdd const& p, pdd const& q, pdd const& r) {
        auto& m = p.manager();

        if (q.is_zero() && p == r)
            return l_true;

        if (q.is_val() && q.val() >= m.power_of_2() && r.is_val())
            return to_lbool(r.is_zero());

        if (p.is_val() && q.is_val() && r.is_val()) {
            SASSERT(q.val().is_unsigned());  // otherwise, previous condition should have been triggered
            // TODO: use right-shift operation instead of division
            auto divisor = rational::power_of_two(q.val().get_unsigned());
            return to_lbool(r.val() == div(p.val(), divisor));
        }

        // TODO: other cases when we know lower bound of q,
        //       e.g, q = 2^k*q1 + q2, where q2 is a constant.
        return l_undef;
    }

    /**
     * Enforce axioms for constraint: r == p << q
     *
     *      q >= K  ->  r = 0
     *      q >= k  ->  r = 0  \/  r >= 2^k
     *      q >= k  ->  r[i] = 0 for i < k
     *      q = k   ->  r[i+k] = p[i] for 0 <= i < K - k
     *      r != 0  ->  r >= p
     *      q = 0   ->  r = p
     */
    clause_ref op_constraint::lemma_shl(solver& s, assignment const& a) {
        auto const pv = a.apply_to(p());
        auto const qv = a.apply_to(q());
        auto const rv = a.apply_to(r());
        unsigned const K = p().manager().power_of_2();

        signed_constraint const shl(this, true);

        if (pv.is_val() && !pv.is_zero() && !pv.is_one() && rv.is_val() && !rv.is_zero() && rv.val() < pv.val())
            // r != 0  ->  r >= p
            // r = 0 \/ r >= p      (equivalent)
            // r-1 >= p-1           (equivalent unit constraint to better support narrowing)
            return s.mk_clause(~shl, s.ule(p() - 1, r() - 1), true);
        else if (qv.is_val() && qv.val() >= K && rv.is_val() && !rv.is_zero())
            // q >= K  ->  r = 0
            return s.mk_clause(~shl, ~s.ule(K, q()), s.eq(r()), true);
        else if (qv.is_zero() && pv.is_val() && rv.is_val() && rv != pv)
            // q = 0  ->  r = p
            return s.mk_clause(~shl, ~s.eq(q()), s.eq(r(), p()), true);
        else if (qv.is_val() && !qv.is_zero() && qv.val() < K && rv.is_val() &&
            !rv.is_zero() && rv.val() < rational::power_of_two(qv.val().get_unsigned()))
            // q >= k  ->  r = 0  \/  r >= 2^k  (intuitive version)
            // q >= k  ->  r - 1 >= 2^k - 1     (equivalent unit constraint to better support narrowing)
            return s.mk_clause(~shl, ~s.ule(qv.val(), q()), s.ule(rational::power_of_two(qv.val().get_unsigned()) - 1, r() - 1), true);
        else if (pv.is_val() && rv.is_val() && qv.is_val() && !qv.is_zero()) {
            unsigned const k = qv.val().get_unsigned();
            // q = k  ->  r[i+k] = p[i] for 0 <= i < K - k
            for (unsigned i = 0; i < K - k; ++i) {
                if (rv.val().get_bit(i + k) && !pv.val().get_bit(i)) {
                    return s.mk_clause(~shl, ~s.eq(q(), k), ~s.bit(r(), i + k), s.bit(p(), i), true);
                }
                if (!rv.val().get_bit(i + k) && pv.val().get_bit(i)) {
                    return s.mk_clause(~shl, ~s.eq(q(), k), s.bit(r(), i + k), ~s.bit(p(), i), true);
                }
            }
        }
        else {
            SASSERT(!(pv.is_val() && qv.is_val() && rv.is_val()));
        }
        return {};
    }

    /** Evaluate constraint: r == p << q */
    lbool op_constraint::eval_shl(pdd const& p, pdd const& q, pdd const& r) {
        auto& m = p.manager();

        if (q.is_zero() && p == r)
            return l_true;

        if (q.is_val() && q.val() >= m.power_of_2() && r.is_val())
            return to_lbool(r.is_zero());

        if (p.is_val() && q.is_val() && r.is_val()) {
            SASSERT(q.val().is_unsigned());  // otherwise, previous condition should have been triggered
            // TODO: use left-shift operation instead of multiplication?
            auto factor = rational::power_of_two(q.val().get_unsigned());
            return to_lbool(r == p * m.mk_val(factor));
        }

        // TODO: other cases when we know lower bound of q,
        //       e.g, q = 2^k*q1 + q2, where q2 is a constant.
        //       (bounds should be tracked by viable, then just use min_viable here)
        return l_undef;
    }

    void op_constraint::activate_and(solver& s) {
        auto x = p(), y = q();
        if (x.is_val())
            std::swap(x, y);
        if (!y.is_val())
            return;
        auto& m = x.manager();
        auto yv = y.val();
        if (!(yv + 1).is_power_of_two())
            return;
        signed_constraint const andc(this, true);
        if (yv == m.max_value()) 
            s.add_clause(~andc, s.eq(x, r()), false);
        else if (yv == 0)
            s.add_clause(~andc, s.eq(r()), false);
        else {
            unsigned K = m.power_of_2();
            unsigned k = yv.get_num_bits();
            SASSERT(k < K);
            rational exp = rational::power_of_two(K - k);
            s.add_clause(~andc, s.eq(x * exp, r() * exp), false);
            s.add_clause(~andc, s.ule(r(), y), false);  // maybe always activate these constraints regardless?
        }        
    }

    /**
     * Produce lemmas for constraint: r == p & q
     * r <= p
     * r <= q
     * p = q => r = p
     * p[i] && q[i] = r[i]
     * p = 2^K - 1 => q = r
     * q = 2^K - 1 => p = r
     * p = 2^k - 1 => r*2^{K - k} = q*2^{K - k}
     * q = 2^k - 1 => r*2^{K - k} = p*2^{K - k}
     * r = 0 && q != 0 & p = 2^k - 1 => q >= 2^k   
     * r = 0 && p != 0 & q = 2^k - 1 => p >= 2^k   
     */
    clause_ref op_constraint::lemma_and(solver& s, assignment const& a) {
        auto& m = p().manager();
        auto pv = a.apply_to(p());
        auto qv = a.apply_to(q());
        auto rv = a.apply_to(r());

        signed_constraint const andc(this, true);

        // r <= p
        if (pv.is_val() && rv.is_val() && rv.val() > pv.val())
            return s.mk_clause(~andc, s.ule(r(), p()), true);
        // r <= q
        if (qv.is_val() && rv.is_val() && rv.val() > qv.val())
            return s.mk_clause(~andc, s.ule(r(), q()), true);
        // p = q => r = p
        if (pv.is_val() && qv.is_val() && rv.is_val() && pv == qv && rv != pv)
            return s.mk_clause(~andc, ~s.eq(p(), q()), s.eq(r(), p()), true);
        if (pv.is_val() && qv.is_val() && rv.is_val()) {
            // p = -1 => r = q
            if (pv.val() == m.max_value() && qv != rv)
                return s.mk_clause(~andc, ~s.eq(p(), m.max_value()), s.eq(q(), r()), true);
            // q = -1 => r = p
            if (qv.val() == m.max_value() && pv != rv)
                return s.mk_clause(~andc, ~s.eq(q(), m.max_value()), s.eq(p(), r()), true);

            unsigned K = m.power_of_2();
            // p = 2^k - 1 => r*2^{K - k} = q*2^{K - k}
            // TODO
            // if ((pv.val() + 1).is_power_of_two() ...)

            // q = 2^k - 1 => r*2^{K - k} = p*2^{K - k}

            // r = 0 && q != 0 & p = 2^k - 1 => q >= 2^k   
            if ((pv.val() + 1).is_power_of_two() && rv.val() > pv.val())
                return s.mk_clause(~andc, ~s.eq(r()), ~s.eq(p(), pv.val()), s.eq(q()), s.ult(p(), q()), true);
            // r = 0 && p != 0 & q = 2^k - 1 => p >= 2^k
            if (rv.is_zero() && (qv.val() + 1).is_power_of_two() && pv.val() <= qv.val())
                return s.mk_clause(~andc, ~s.eq(r()), ~s.eq(q(), qv.val()), s.eq(p()),s.ult(q(), p()), true);
            
            for (unsigned i = 0; i < K; ++i) {
                bool pb = pv.val().get_bit(i);
                bool qb = qv.val().get_bit(i);
                bool rb = rv.val().get_bit(i);
                if (rb == (pb && qb))
                    continue;
                if (pb && qb && !rb)
                    return s.mk_clause(~andc, ~s.bit(p(), i), ~s.bit(q(), i), s.bit(r(), i), true);
                else if (!pb && rb)
                    return s.mk_clause(~andc, s.bit(p(), i), ~s.bit(r(), i), true);
                else if (!qb && rb)
                    return s.mk_clause(~andc, s.bit(q(), i), ~s.bit(r(), i), true);
                else
                    UNREACHABLE();
                return {};
            }
        }
        // Propagate r if p or q are 0
        if (pv.is_zero() && !rv.is_zero())  // rv not necessarily fully evaluated
            return s.mk_clause(~andc, s.ule(r(), p()), true);
        if (qv.is_zero() && !rv.is_zero())  // rv not necessarily fully evaluated
            return s.mk_clause(~andc, s.ule(r(), q()), true);
        return {};
    }

    /** Evaluate constraint: r == p & q */
    lbool op_constraint::eval_and(pdd const& p, pdd const& q, pdd const& r) {
        if ((p.is_zero() || q.is_zero()) && r.is_zero())
            return l_true;

        if (p.is_val() && q.is_val() && r.is_val())
            return r.val() == bitwise_and(p.val(), q.val()) ? l_true : l_false;

        return l_undef;
    }

    void op_constraint::add_to_univariate_solver(solver& s, univariate_solver& us, unsigned dep, bool is_positive) const {
        auto p_coeff = s.subst(p()).get_univariate_coefficients();
        auto q_coeff = s.subst(q()).get_univariate_coefficients();
        auto r_coeff = s.subst(r()).get_univariate_coefficients();
        switch (m_op) {
        case code::lshr_op:
            us.add_lshr(p_coeff, q_coeff, r_coeff, !is_positive, dep);
            break;
        case code::shl_op:
            us.add_shl(p_coeff, q_coeff, r_coeff, !is_positive, dep);
            break;
        case code::and_op:
            us.add_and(p_coeff, q_coeff, r_coeff, !is_positive, dep);
            break;
        default:
            NOT_IMPLEMENTED_YET();
            break;
        }
    }

}