{-|
== 2-satisfiability

Solves the 2-sat problem, which assigns boolean values to variables (a, b, c...) such that it
satisfies a boolean expression which is a conjunction of clauses, where each clause is a disjunction
of two variables. For example, (a || b) && (a || not c) && (not b || not c).
The solution is obtained from the strongly connected components of the implication graph constructed
from the clauses.

Sources:

* https://en.wikipedia.org/wiki/2-satisfiability
* https://cp-algorithms.com/graph/2SAT.html

-}

module TwoSat
    ( Var(..)
    , solve2Sat
    ) where

import Control.DeepSeq
import Data.Array
import Data.Bits
import Data.Foldable
import Data.Graph

data Var = Id !Int | Not !Int deriving Int -> Var -> ShowS
[Var] -> ShowS
Var -> String
(Int -> Var -> ShowS)
-> (Var -> String) -> ([Var] -> ShowS) -> Show Var
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Var] -> ShowS
$cshowList :: [Var] -> ShowS
show :: Var -> String
$cshow :: Var -> String
showsPrec :: Int -> Var -> ShowS
$cshowsPrec :: Int -> Var -> ShowS
Show

-- | Solves the 2-sat problem where each variable is represented by an Int in the range (l, r). Returns a
-- list of assignments for the variables l to r, in order. O(n + m) where n = r - l + 1 and
-- m = length xys.
solve2Sat :: (Int, Int) -> [(Var, Var)] -> Maybe [Bool]
solve2Sat :: (Int, Int) -> [(Var, Var)] -> Maybe [Bool]
solve2Sat (Int
l, Int
r) [(Var, Var)]
xys = [Int] -> Maybe [Bool]
forall a. Ord a => [a] -> Maybe [Bool]
assign ([Int] -> Maybe [Bool]) -> [Int] -> Maybe [Bool]
forall a b. (a -> b) -> a -> b
$ Array Int Int -> [Int]
forall i e. Array i e -> [e]
elems Array Int Int
comp where
    n :: Int
n = Int
r Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
    bnds :: (Int, Int)
bnds = (Int
0, Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
    idx :: Var -> Int
idx (Id Int
x)  = (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l) Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
2
    idx (Not Int
x) = (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
l) Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
    g :: Graph
g = (Int, Int) -> [(Int, Int)] -> Graph
buildG (Int, Int)
bnds ([(Int, Int)] -> Graph) -> [(Int, Int)] -> Graph
forall a b. (a -> b) -> a -> b
$ do
        (Var
x, Var
y) <- [(Var, Var)]
xys
        [(Var -> Int
idx Var
x Int -> Int -> Int
forall a. Bits a => a -> a -> a
`xor` Int
1, Var -> Int
idx Var
y), (Var -> Int
idx Var
y Int -> Int -> Int
forall a. Bits a => a -> a -> a
`xor` Int
1, Var -> Int
idx Var
x)]
    comp :: Array Int Int
comp = (Int, Int) -> [(Int, Int)] -> Array Int Int
forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array (Int, Int)
bnds ([(Int, Int)] -> Array Int Int) -> [(Int, Int)] -> Array Int Int
forall a b. (a -> b) -> a -> b
$ [[(Int, Int)]] -> [(Int, Int)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[(Int, Int)]] -> [(Int, Int)]) -> [[(Int, Int)]] -> [(Int, Int)]
forall a b. (a -> b) -> a -> b
$ ([Int] -> [Int] -> [(Int, Int)])
-> [[Int]] -> [[Int]] -> [[(Int, Int)]]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith [Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip ((Tree Int -> [Int]) -> [Tree Int] -> [[Int]]
forall a b. (a -> b) -> [a] -> [b]
map Tree Int -> [Int]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList ([Tree Int] -> [[Int]]) -> [Tree Int] -> [[Int]]
forall a b. (a -> b) -> a -> b
$ Graph -> [Tree Int]
scc Graph
g) ((Int -> [Int]) -> [Int] -> [[Int]]
forall a b. (a -> b) -> [a] -> [b]
map Int -> [Int]
forall a. a -> [a]
repeat [Int
1 :: Int ..])
    assign :: [a] -> Maybe [Bool]
assign [] = [Bool] -> Maybe [Bool]
forall a. a -> Maybe a
Just []
    assign ~(a
cx:a
cnotx:[a]
rest)
        | a
cx a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
cnotx = Maybe [Bool]
forall a. Maybe a
Nothing
        | Bool
otherwise   = ((a
cx a -> a -> Bool
forall a. Ord a => a -> a -> Bool
< a
cnotx)Bool -> [Bool] -> [Bool]
forall a. a -> [a] -> [a]
:) ([Bool] -> [Bool]) -> Maybe [Bool] -> Maybe [Bool]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [a] -> Maybe [Bool]
assign [a]
rest

--------------------------------------------------------------------------------
-- For tests

instance NFData Var where
    rnf :: Var -> ()
rnf (Id Int
x)  = Int -> ()
forall a. NFData a => a -> ()
rnf Int
x
    rnf (Not Int
x) = Int -> ()
forall a. NFData a => a -> ()
rnf Int
x