DRMacIver's Notebook

from pysat.solvers import Glucose4
from .cnfbuilder import CNFBuilder


def setminus(A: list[list[int]], B: list[list[int]]) -> list[list[int]]:
    """Returns a set of SAT clauses such that the satisfying
    solutions are precisely the satisfying solutions of A that
    do not satisfy B.

    A, B are SAT formulae represented in CNF in the usual way
    (positive literals correspond to a variable, negative literals
    to the negation fo that variable). The returned formula is in
    the same format.
    """

    # Set of clauses that we incrementally expand until their
    # satisfying solutions are solutions(A) - solutions(B).
    result_clauses = list(A)

    B = list(B)

    # We are only interested in clauses in B that are not also in A,
    # as every clause that is in A is always satisfied by an element
    # of our result.
    target_clauses = {tuple(sorted(c)) for c in B} - {tuple(sorted(c)) for c in A}

    # Every element of A satisfies B, so the set is empty.
    if not target_clauses:
        return [[1], [-1]]

    builder = CNFBuilder(A)

    b_is_satisfied = builder.all([builder.any(c) for c in target_clauses])

    # This is a set of clauses for the intersection of our current
    # best answer with B. If these are satisfiable we're not done yet.
    intersection_clauses = list(A) + list(B)

    # This is a set of clauses with new variables introduced
    # such that an assignment is valid iff it can be extended
    # to a valid assignment of these clauses.
    general_clauses = builder.clauses + [[-b_is_satisfied]]

    while True:
        solver = Glucose4(intersection_clauses)
        if not solver.solve():
            return result_clauses

        # This is a solution to our result_clauses that is also in
        # B, and so indicates that we're not done.
        counter_example = solver.get_model()

        # We now feed it in to the general clauses. This cannot be
        # extended to a satisfying solution, because a solution is
        # satisfying for general_clauses if and only if it corresponds
        # to the unique extension of some satisfaction of A setminus B.
        clause_generator = Glucose4(general_clauses)
        solved = clause_generator.solve(assumptions=counter_example)
        assert not solved

        # This gives us a minimal clause that any solution to A setminus B must
        # satisfy k
        core = clause_generator.get_core()
        assert core is not None
        new_clause = [-v for v in core]
        general_clauses.append(new_clause)
        intersection_clauses.append(new_clause)
        result_clauses.append(new_clause)

class CNFBuilder:
    def __init__(self, clauses):
        self.clauses = list(clauses)
        self.__nextvar = (
            max(abs(literal) for clause in clauses for literal in clause) + 1
        )

    def any(self, variables) -> int:
        result = self.__newvar()
        # Any variable being true implies the result variable is true.
        for v in variables:
            self.__implies(v, result)

        # If none of the variables are true, then any(variables) is false.
        self.clauses.append([-result] + [v for v in variables])
        return result

    def all(self, variables) -> int:
        result = self.__newvar()
        for v in variables:
            self.__implies(result, v)

        # If all of the variables are true, then all(variables) is true.
        self.clauses.append([result] + [-v for v in variables])
        return result

    def __implies(self, l, r):
        self.clauses.append([-l, r])

    def __newvar(self):
        result = self.__nextvar
        self.__nextvar += 1
        return result

def remove_one_variable(clauses, variable):
    """Return a set of SAT clauses not containing `variable` that accepts only
    assignments that can be turned into satisfying assignments of `clauses` through
    a suitable choice of assignment for `variable`."""
    assert variable > 0

    solver = Glucose4(clauses)

    # Current set of clauses representing our result set. Initially starts
    # out as all clauses that do not reference the variable to remove.
    refined_clauses: list[Clause] = []

    # All clauses in the original where the variable appears positively.
    positive_clauses = []

    # All clauses in the original where the variable appears negatively.
    negative_clauses = []

    for clause in clauses:
        if variable in clause:
            clause = sorted(set(clause) - {variable})
            positive_clauses.append(clause)
        elif -variable in clause:
            clause = sorted(set(clause) - {-variable})
            negative_clauses.append(clause)
        else:
            refined_clauses.append(clause)

    assert positive_clauses or negative_clauses

    builder = CNFBuilder(refined_clauses)

    positive_satisfied = builder.all(
        [builder.any(clause) for clause in positive_clauses]
    )

    negative_satisfied = builder.all(
        [builder.any(clause) for clause in negative_clauses]
    )

    # These clauses are satisfied if and only if all of our
    # current result clauses are satisfied but there is no possible
    # assignment of the removed variable that makes the original set
    # of clauses satisfied.
    oracle_clauses = builder.clauses
    oracle_clauses.append([-positive_satisfied])
    oracle_clauses.append([-negative_satisfied])
    oracle = Glucose4(oracle_clauses)

    while True:
        # Every solution of our refined_clauses has a satisfying extension,
        # so we're done.
        if not oracle.solve():
            return refined_clauses

        # We've got some solution to the refined clauses that cannot be extended.
        # Turn it into a set of assumptions that cannot be satisfied in the original
        # SAT problem.
        counterexample = oracle.get_model()
        assert counterexample is not None
        assumptions = [
            literal for literal in counterexample if abs(literal) != variable
        ]

        solved_assumptions = solver.solve(assumptions=assumptions)
        assert not solved_assumptions
        core = solver.get_core()
        assert core is not None

        # We've learned a new clause that must always hold for the remaining variables.
        new_clause = tuple([-literal for literal in core])
        refined_clauses.append(new_clause)
        oracle.add_clause(new_clause)