CutEnumSet.hs

module GraphColoring.CutEnumSet where

import qualified Math.SetCover.Exact as ESC

import qualified Data.EnumBitSet as EnumBitSet
import qualified Data.EnumSet as EnumSet
import qualified Data.List as List
import Data.Monoid (First(First, getFirst))
import Data.EnumSet (EnumSet)
import Data.Ord.HT (comparing)
import Data.Either.HT (maybeRight)
import Data.Maybe.HT (toMaybe)
import Data.Maybe (mapMaybe)
import Data.Word (Word8)
import Text.Printf (printf)

{- $setup
>>> import qualified Math.SetCover.Exact as ESC
>>>
>>> import qualified Data.EnumSet as EnumSet
>>> import qualified Data.List.Key as Key
>>> import Data.Word (Word32)
>>>
>>> import Control.Applicative (liftA2)
>>>
>>> import qualified Test.QuickCheck as QC
>>> import Test.QuickCheck ((===))
-}


type Label node color = Either (color, color) (node, color)
type Assign node color set = ESC.Assign (Label node color) set

{- |
Re-implementation of SetCover.Exact functions with an extended state,
that also maintains the already used colors in an efficient 'EnumSet'.
-}
data State node color set =
   State {
      stateSetCover :: ESC.State (Label node color) set,
      stateUsedColors :: EnumSet color
   }

initState :: (Enum color, ESC.Set set) =>
   [Assign node color set] -> State node color set
initState asns =
   State {
      stateSetCover = ESC.initState asns,
      stateUsedColors = EnumSet.empty
   }

{-# INLINE updateState #-}
updateState :: (Enum color, ESC.Set set) =>
   Assign node color set ->
   State node color set -> State node color set
updateState asn s =
   State {
      stateSetCover = ESC.updateState asn $ stateSetCover s,
      stateUsedColors =
         stateUsedColors s <>
         case ESC.label asn of
            Right (_,color) -> EnumSet.singleton color
            Left (color0,color1) ->
               EnumSet.singleton color0 <> EnumSet.singleton color1
   }

{- |
Without loss of generality, when new colors are choosen,
only try the lexicographically lowest single color or color pair.
-}
uniqueNewColors :: (Enum color) =>
   EnumSet color -> [Assign node color set] -> [Assign node color set]
uniqueNewColors alreadyUsedColors =
   (\(ordinaryAttempts,
         tryNewColors0, tryNewColors1, tryNewColors2, tryNewColors3) ->
      ordinaryAttempts ++
      mapMaybe getFirst
         [tryNewColors0, tryNewColors1, tryNewColors2, tryNewColors3])
   .
   foldMap
      (\asn ->
         let single = First . Just in
         case ESC.label asn of
            Left (color0,color1) ->
               case (EnumSet.member color0 alreadyUsedColors,
                     EnumSet.member color1 alreadyUsedColors) of
                  (True, True) -> ([asn], mempty, mempty, mempty, mempty)
                  (False, False) ->
                     (mempty, mempty, single asn, mempty, mempty)
                  (False, True) ->
                     (mempty, mempty, mempty, single asn, mempty)
                  (True, False) ->
                     (mempty, mempty, mempty, mempty, single asn)
            Right (_node,color) ->
               if EnumSet.member color alreadyUsedColors
                  then ([asn], mempty, mempty, mempty, mempty)
                  else (mempty, single asn, mempty, mempty, mempty))



data UsagePattern = NewNode | NewEdge0 | NewEdge1 | NewEdge01
   deriving (Eq, Ord, Enum, Show)

{- |
This is the same as uniqueNewColors
but performs a single scan through the list
and bookkeeps already seen patterns of new colors.
Order of assignments is maintained.


prop> :{
   QC.forAll (fmap (\n -> take (1 + mod n 5) ['a'..]) QC.arbitrary) $ \set ->
   QC.forAll (QC.sublistOf set) $ \used ->
   let usedSet = EnumSet.fromList used in
   let genLabel =
         QC.oneof
            [fmap Left $ liftA2 (,) (QC.elements set) (QC.elements set),
             fmap Right $
               liftA2 (,) (QC.arbitrary :: QC.Gen Integer) (QC.elements set)]
       genAssignSet :: QC.Gen Word32
       genAssignSet = QC.arbitrary
   in
   QC.forAll (QC.listOf (liftA2 ESC.assign genLabel genAssignSet)) $ \assigns ->
   let normalize =
         Key.sort fst . map (\(ESC.Assign label aset) -> (label,aset)) in

   normalize (uniqueNewColors usedSet assigns)
   ===
   normalize (uniqueNewColorsStable usedSet assigns)
:}
-}
uniqueNewColorsStable :: (Enum color) =>
   EnumSet color -> [Assign node color set] -> [Assign node color set]
uniqueNewColorsStable alreadyUsedColors asns =
   map fst $ filter snd $ zip asns $
   snd $
   List.mapAccumL
      (\seen mpat ->
         case mpat of
            Nothing -> (seen, True)
            Just pat ->
               (EnumBitSet.set pat seen, not (EnumBitSet.get pat seen)))
      (EnumBitSet.empty :: EnumBitSet.T Word8 UsagePattern)
   $
   flip map asns
      (\asn ->
         case ESC.label asn of
            Left (color0,color1) ->
               case (EnumSet.member color0 alreadyUsedColors,
                     EnumSet.member color1 alreadyUsedColors) of
                  (True, True) -> Nothing
                  (False, False) -> Just NewEdge01
                  (False, True) -> Just NewEdge0
                  (True, False) -> Just NewEdge1
            Right (_node,color) ->
               toMaybe (EnumSet.notMember color alreadyUsedColors) NewNode)

{-# INLINE step #-}
step :: (Enum color, ESC.Set set) =>
   State node color set -> [State node color set]
step s =
   map (flip updateState s) $
   uniqueNewColorsStable (stateUsedColors s) $
   let se = stateSetCover s in
   ESC.minimize (ESC.freeElements se) (ESC.availableSubsets se)

{-# INLINE search #-}
search :: (Enum color, ESC.Set set) =>
   State node color set -> [[Label node color]]
search s =
   let se = stateSetCover s in
   if ESC.null (ESC.freeElements se)
     then [ESC.usedSubsets se]
     else step s >>= search

run :: (ESC.Set set) => [Assign Int Char set] -> IO ()
run asns =
   case search $ initState asns of
      [] -> putStrLn "coloring impossible"
      solution:_ ->
         mapM_ (\(node,color) -> printf "Node %d: %c\n" node color) $
         List.sortBy (comparing fst) $ mapMaybe maybeRight solution