{-# language FlexibleInstances, MultiParamTypeClasses, FlexibleContexts #-}
{-# language TupleSections #-}

module Satchmo.Set.Data

( Set , unknown, unknownSingleton, constant
, member, keys, keysSet, keys, assocs, elems
, all2, common2
) 

where

import Satchmo.Code
import qualified Satchmo.Boolean as B

import Satchmo.SAT

import qualified Data.Set as S
import qualified Data.Map.Strict as M

import Control.Monad ( guard, forM )
import Control.Applicative ( (<$>), (<*>) )
import Data.List ( tails )

newtype Set a = Set (M.Map a B.Boolean)

instance ( Functor m, Decode m B.Boolean Bool, Ord a )
         => Decode m (Set a) ( S.Set a) where
    decode :: Set a -> m (Set a)
decode (Set Map a Boolean
m) = 
        Map a Bool -> Set a
forall k a. Map k a -> Set k
M.keysSet (Map a Bool -> Set a)
-> (Map a Bool -> Map a Bool) -> Map a Bool -> Set a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Bool -> Bool) -> Map a Bool -> Map a Bool
forall a k. (a -> Bool) -> Map k a -> Map k a
M.filter Bool -> Bool
forall a. a -> a
id (Map a Bool -> Set a) -> m (Map a Bool) -> m (Set a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map a Boolean -> m (Map a Bool)
forall (m :: * -> *) c a. Decode m c a => c -> m a
decode Map a Boolean
m

keys :: Set k -> [k]
keys (Set Map k Boolean
m) = Map k Boolean -> [k]
forall k a. Map k a -> [k]
M.keys Map k Boolean
m
keysSet :: Set k -> Set k
keysSet (Set Map k Boolean
m) = Map k Boolean -> Set k
forall k a. Map k a -> Set k
M.keysSet Map k Boolean
m
assocs :: Set k -> [(k, Boolean)]
assocs (Set Map k Boolean
m) = Map k Boolean -> [(k, Boolean)]
forall k a. Map k a -> [(k, a)]
M.assocs Map k Boolean
m
elems :: Set k -> [Boolean]
elems (Set Map k Boolean
m) = Map k Boolean -> [Boolean]
forall k a. Map k a -> [a]
M.elems Map k Boolean
m

member :: k -> Set k -> m Boolean
member k
x (Set Map k Boolean
m) = case k -> Map k Boolean -> Maybe Boolean
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
x Map k Boolean
m of
    Maybe Boolean
Nothing -> Bool -> m Boolean
forall (m :: * -> *). MonadSAT m => Bool -> m Boolean
B.constant Bool
False
    Just Boolean
y  -> Boolean -> m Boolean
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Boolean
y


-- | allocate an unknown subset of these elements
unknown :: ( B.MonadSAT m , Ord a )
         => [a] -> m (Set a)
unknown :: forall (m :: * -> *) a. (MonadSAT m, Ord a) => [a] -> m (Set a)
unknown [a]
xs = Map a Boolean -> Set a
forall a. Map a Boolean -> Set a
Set (Map a Boolean -> Set a)
-> ([(a, Boolean)] -> Map a Boolean) -> [(a, Boolean)] -> Set a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(a, Boolean)] -> Map a Boolean
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList 
     ([(a, Boolean)] -> Set a) -> m [(a, Boolean)] -> m (Set a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ( [a] -> (a -> m (a, Boolean)) -> m [(a, Boolean)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [a]
xs ((a -> m (a, Boolean)) -> m [(a, Boolean)])
-> (a -> m (a, Boolean)) -> m [(a, Boolean)]
forall a b. (a -> b) -> a -> b
$ \ a
x -> (a
x,) (Boolean -> (a, Boolean)) -> m Boolean -> m (a, Boolean)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m Boolean
forall (m :: * -> *). MonadSAT m => m Boolean
B.boolean )

unknownSingleton :: [k] -> m (Set k)
unknownSingleton [k]
xs = do
    Set k
s <- [k] -> m (Set k)
forall (m :: * -> *) a. (MonadSAT m, Ord a) => [a] -> m (Set a)
unknown [k]
xs
    [Boolean] -> m ()
forall (m :: * -> *). MonadSAT m => [Boolean] -> m ()
B.assert ([Boolean] -> m ()) -> [Boolean] -> m ()
forall a b. (a -> b) -> a -> b
$ Set k -> [Boolean]
forall {k}. Set k -> [Boolean]
elems Set k
s
    [m ()] -> m ()
forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, Monad m) =>
t (m a) -> m ()
sequence_ ([m ()] -> m ()) -> [m ()] -> m ()
forall a b. (a -> b) -> a -> b
$ do 
       Boolean
x : [Boolean]
ys <- [Boolean] -> [[Boolean]]
forall a. [a] -> [[a]]
tails ([Boolean] -> [[Boolean]]) -> [Boolean] -> [[Boolean]]
forall a b. (a -> b) -> a -> b
$ Set k -> [Boolean]
forall {k}. Set k -> [Boolean]
elems Set k
s ; Boolean
y <- [Boolean]
ys
       m () -> [m ()]
forall a. a -> [a]
forall (m :: * -> *) a. Monad m => a -> m a
return (m () -> [m ()]) -> m () -> [m ()]
forall a b. (a -> b) -> a -> b
$ [Boolean] -> m ()
forall (m :: * -> *). MonadSAT m => [Boolean] -> m ()
B.assert [ Boolean -> Boolean
B.not Boolean
x, Boolean -> Boolean
B.not Boolean
y ]
    Set k -> m (Set k)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Set k
s

constant :: ( B.MonadSAT m , Ord a )
         => [a] -> m (Set a)
constant :: forall (m :: * -> *) a. (MonadSAT m, Ord a) => [a] -> m (Set a)
constant [a]
xs = Map a Boolean -> Set a
forall a. Map a Boolean -> Set a
Set (Map a Boolean -> Set a)
-> ([(a, Boolean)] -> Map a Boolean) -> [(a, Boolean)] -> Set a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(a, Boolean)] -> Map a Boolean
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList 
     ([(a, Boolean)] -> Set a) -> m [(a, Boolean)] -> m (Set a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ( [a] -> (a -> m (a, Boolean)) -> m [(a, Boolean)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [a]
xs ((a -> m (a, Boolean)) -> m [(a, Boolean)])
-> (a -> m (a, Boolean)) -> m [(a, Boolean)]
forall a b. (a -> b) -> a -> b
$ \ a
x -> (a
x,) (Boolean -> (a, Boolean)) -> m Boolean -> m (a, Boolean)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bool -> m Boolean
forall (m :: * -> *). MonadSAT m => Bool -> m Boolean
B.constant Bool
True )

all2 :: (Boolean -> Boolean -> m Boolean) -> Set k -> Set k -> m Boolean
all2 Boolean -> Boolean -> m Boolean
f Set k
s Set k
t = [Boolean] -> m Boolean
forall (m :: * -> *). MonadSAT m => [Boolean] -> m Boolean
B.and
 ([Boolean] -> m Boolean) -> m [Boolean] -> m Boolean
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< [k] -> (k -> m Boolean) -> m [Boolean]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ( Set k -> [k]
forall a. Set a -> [a]
S.toList (Set k -> [k]) -> Set k -> [k]
forall a b. (a -> b) -> a -> b
$ Set k -> Set k -> Set k
forall a. Ord a => Set a -> Set a -> Set a
S.union (Set k -> Set k
forall {k}. Set k -> Set k
keysSet Set k
s)(Set k -> Set k
forall {k}. Set k -> Set k
keysSet Set k
t))
 ( \ k
x -> do Boolean
a <- k -> Set k -> m Boolean
forall {k} {m :: * -> *}.
(Ord k, MonadSAT m) =>
k -> Set k -> m Boolean
member k
x Set k
s; Boolean
b <- k -> Set k -> m Boolean
forall {k} {m :: * -> *}.
(Ord k, MonadSAT m) =>
k -> Set k -> m Boolean
member k
x Set k
t; Boolean -> Boolean -> m Boolean
f Boolean
a Boolean
b )

common2 :: (Boolean -> Boolean -> f Boolean) -> Set a -> Set a -> f (Set a)
common2 Boolean -> Boolean -> f Boolean
f Set a
s Set a
t = Map a Boolean -> Set a
forall a. Map a Boolean -> Set a
Set (Map a Boolean -> Set a)
-> ([(a, Boolean)] -> Map a Boolean) -> [(a, Boolean)] -> Set a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(a, Boolean)] -> Map a Boolean
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(a, Boolean)] -> Set a) -> f [(a, Boolean)] -> f (Set a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
 [a] -> (a -> f (a, Boolean)) -> f [(a, Boolean)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM ( Set a -> [a]
forall a. Set a -> [a]
S.toList (Set a -> [a]) -> Set a -> [a]
forall a b. (a -> b) -> a -> b
$ Set a -> Set a -> Set a
forall a. Ord a => Set a -> Set a -> Set a
S.union (Set a -> Set a
forall {k}. Set k -> Set k
keysSet Set a
s)(Set a -> Set a
forall {k}. Set k -> Set k
keysSet Set a
t))
 ( \ a
x -> do Boolean
a <- a -> Set a -> f Boolean
forall {k} {m :: * -> *}.
(Ord k, MonadSAT m) =>
k -> Set k -> m Boolean
member a
x Set a
s; Boolean
b <- a -> Set a -> f Boolean
forall {k} {m :: * -> *}.
(Ord k, MonadSAT m) =>
k -> Set k -> m Boolean
member a
x Set a
t
             Boolean
y <- Boolean -> Boolean -> f Boolean
f Boolean
a Boolean
b ; (a, Boolean) -> f (a, Boolean)
forall a. a -> f a
forall (m :: * -> *) a. Monad m => a -> m a
return (a
x,Boolean
y) )