module Math.SetCover.EnumMap where

import qualified Math.SetCover.BitPosition as BitPos
import qualified Math.SetCover.BitSet as BitSet

import qualified Data.EnumMap as EnumMap; import Data.EnumMap (EnumMap)
import qualified Data.EnumSet as EnumSet; import Data.EnumSet (EnumSet)
import qualified Data.IntMap as IntMap; import Data.IntMap (IntMap)
import qualified Data.IntSet as IntSet; import Data.IntSet (IntSet)
import qualified Data.Map as Map
import qualified Data.Set as Set

import Prelude hiding (const)


-- EnumMap.fromSet is available from containers-0.5
const :: (Enum e) => a -> EnumSet e -> EnumMap e a
const :: forall e a. Enum e => a -> EnumSet e -> EnumMap e a
const a
a = [(e, a)] -> EnumMap e a
forall k a. Enum k => [(k, a)] -> EnumMap k a
EnumMap.fromAscList ([(e, a)] -> EnumMap e a)
-> (EnumSet e -> [(e, a)]) -> EnumSet e -> EnumMap e a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (e -> (e, a)) -> [e] -> [(e, a)]
forall a b. (a -> b) -> [a] -> [b]
map (\e
k -> (e
k, a
a)) ([e] -> [(e, a)]) -> (EnumSet e -> [e]) -> EnumSet e -> [(e, a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. EnumSet e -> [e]
forall k. Enum k => EnumSet k -> [k]
EnumSet.toAscList

intersection :: (Enum e) => EnumMap e a -> EnumSet e -> EnumMap e a
intersection :: forall e a. Enum e => EnumMap e a -> EnumSet e -> EnumMap e a
intersection EnumMap e a
m EnumSet e
s = EnumMap e a -> EnumMap e () -> EnumMap e a
forall k a b. EnumMap k a -> EnumMap k b -> EnumMap k a
EnumMap.intersection EnumMap e a
m (EnumMap e () -> EnumMap e a) -> EnumMap e () -> EnumMap e a
forall a b. (a -> b) -> a -> b
$ () -> EnumSet e -> EnumMap e ()
forall e a. Enum e => a -> EnumSet e -> EnumMap e a
const () EnumSet e
s

partition ::
   (Enum e) => EnumMap e a -> EnumSet e -> (EnumMap e a, EnumMap e a)
partition :: forall e a.
Enum e =>
EnumMap e a -> EnumSet e -> (EnumMap e a, EnumMap e a)
partition EnumMap e a
m EnumSet e
s =
   let section :: EnumMap e a
section = EnumMap e a -> EnumSet e -> EnumMap e a
forall e a. Enum e => EnumMap e a -> EnumSet e -> EnumMap e a
intersection EnumMap e a
m EnumSet e
s
   in  (EnumMap e a
section, EnumMap e a -> EnumMap e a -> EnumMap e a
forall k a b. EnumMap k a -> EnumMap k b -> EnumMap k a
EnumMap.difference EnumMap e a
m EnumMap e a
section)


{-# INLINE attach #-}
attach :: b -> [a] -> [(a, b)]
attach :: forall b a. b -> [a] -> [(a, b)]
attach b
a = (a -> (a, b)) -> [a] -> [(a, b)]
forall a b. (a -> b) -> [a] -> [b]
map ((a -> b -> (a, b)) -> b -> a -> (a, b)
forall a b c. (a -> b -> c) -> b -> a -> c
flip (,) b
a)

-- Map.fromSet is available from containers-0.5
constMap :: (Ord a) => b -> Set.Set a -> Map.Map a b
constMap :: forall a b. Ord a => b -> Set a -> Map a b
constMap b
a = [(a, b)] -> Map a b
forall k a. Eq k => [(k, a)] -> Map k a
Map.fromAscList ([(a, b)] -> Map a b) -> (Set a -> [(a, b)]) -> Set a -> Map a b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> [a] -> [(a, b)]
forall b a. b -> [a] -> [(a, b)]
attach b
a ([a] -> [(a, b)]) -> (Set a -> [a]) -> Set a -> [(a, b)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set a -> [a]
forall a. Set a -> [a]
Set.toAscList

transposeSet ::
   (Enum e, Ord a) => EnumMap e (Set.Set a) -> Map.Map a (EnumSet e)
transposeSet :: forall e a.
(Enum e, Ord a) =>
EnumMap e (Set a) -> Map a (EnumSet e)
transposeSet =
   (EnumSet e -> EnumSet e -> EnumSet e)
-> [Map a (EnumSet e)] -> Map a (EnumSet e)
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
(a -> a -> a) -> f (Map k a) -> Map k a
Map.unionsWith EnumSet e -> EnumSet e -> EnumSet e
forall k. EnumSet k -> EnumSet k -> EnumSet k
EnumSet.union ([Map a (EnumSet e)] -> Map a (EnumSet e))
-> (EnumMap e (Set a) -> [Map a (EnumSet e)])
-> EnumMap e (Set a)
-> Map a (EnumSet e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. EnumMap e (Map a (EnumSet e)) -> [Map a (EnumSet e)]
forall k a. EnumMap k a -> [a]
EnumMap.elems (EnumMap e (Map a (EnumSet e)) -> [Map a (EnumSet e)])
-> (EnumMap e (Set a) -> EnumMap e (Map a (EnumSet e)))
-> EnumMap e (Set a)
-> [Map a (EnumSet e)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
   (e -> Set a -> Map a (EnumSet e))
-> EnumMap e (Set a) -> EnumMap e (Map a (EnumSet e))
forall k a b. Enum k => (k -> a -> b) -> EnumMap k a -> EnumMap k b
EnumMap.mapWithKey (EnumSet e -> Set a -> Map a (EnumSet e)
forall a b. Ord a => b -> Set a -> Map a b
constMap (EnumSet e -> Set a -> Map a (EnumSet e))
-> (e -> EnumSet e) -> e -> Set a -> Map a (EnumSet e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. e -> EnumSet e
forall k. Enum k => k -> EnumSet k
EnumSet.singleton)


transposeMap ::
   (Enum e, Ord a) => EnumMap e (Map.Map a b) -> Map.Map a (EnumMap e b)
transposeMap :: forall e a b.
(Enum e, Ord a) =>
EnumMap e (Map a b) -> Map a (EnumMap e b)
transposeMap =
   (EnumMap e b -> EnumMap e b -> EnumMap e b)
-> [Map a (EnumMap e b)] -> Map a (EnumMap e b)
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
(a -> a -> a) -> f (Map k a) -> Map k a
Map.unionsWith EnumMap e b -> EnumMap e b -> EnumMap e b
forall k a. EnumMap k a -> EnumMap k a -> EnumMap k a
EnumMap.union ([Map a (EnumMap e b)] -> Map a (EnumMap e b))
-> (EnumMap e (Map a b) -> [Map a (EnumMap e b)])
-> EnumMap e (Map a b)
-> Map a (EnumMap e b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. EnumMap e (Map a (EnumMap e b)) -> [Map a (EnumMap e b)]
forall k a. EnumMap k a -> [a]
EnumMap.elems (EnumMap e (Map a (EnumMap e b)) -> [Map a (EnumMap e b)])
-> (EnumMap e (Map a b) -> EnumMap e (Map a (EnumMap e b)))
-> EnumMap e (Map a b)
-> [Map a (EnumMap e b)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
   (e -> Map a b -> Map a (EnumMap e b))
-> EnumMap e (Map a b) -> EnumMap e (Map a (EnumMap e b))
forall k a b. Enum k => (k -> a -> b) -> EnumMap k a -> EnumMap k b
EnumMap.mapWithKey ((b -> EnumMap e b) -> Map a b -> Map a (EnumMap e b)
forall a b. (a -> b) -> Map a a -> Map a b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((b -> EnumMap e b) -> Map a b -> Map a (EnumMap e b))
-> (e -> b -> EnumMap e b) -> e -> Map a b -> Map a (EnumMap e b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. e -> b -> EnumMap e b
forall k a. Enum k => k -> a -> EnumMap k a
EnumMap.singleton)


constIntMapFromBits :: (BitPos.C bits) => b -> BitSet.Set bits -> IntMap b
constIntMapFromBits :: forall bits b. C bits => b -> Set bits -> IntMap b
constIntMapFromBits b
a = [(Key, b)] -> IntMap b
forall a. [(Key, a)] -> IntMap a
IntMap.fromAscList ([(Key, b)] -> IntMap b)
-> (Set bits -> [(Key, b)]) -> Set bits -> IntMap b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> [Key] -> [(Key, b)]
forall b a. b -> [a] -> [(a, b)]
attach b
a ([Key] -> [(Key, b)])
-> (Set bits -> [Key]) -> Set bits -> [(Key, b)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set bits -> [Key]
forall bits. C bits => Set bits -> [Key]
BitPos.unpack

transposeBitSet ::
   (BitPos.C bits, Enum e) => EnumMap e (BitSet.Set bits) -> IntMap (EnumSet e)
transposeBitSet :: forall bits e.
(C bits, Enum e) =>
EnumMap e (Set bits) -> IntMap (EnumSet e)
transposeBitSet =
   (EnumSet e -> EnumSet e -> EnumSet e)
-> [IntMap (EnumSet e)] -> IntMap (EnumSet e)
forall (f :: * -> *) a.
Foldable f =>
(a -> a -> a) -> f (IntMap a) -> IntMap a
IntMap.unionsWith EnumSet e -> EnumSet e -> EnumSet e
forall k. EnumSet k -> EnumSet k -> EnumSet k
EnumSet.union ([IntMap (EnumSet e)] -> IntMap (EnumSet e))
-> (EnumMap e (Set bits) -> [IntMap (EnumSet e)])
-> EnumMap e (Set bits)
-> IntMap (EnumSet e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. EnumMap e (IntMap (EnumSet e)) -> [IntMap (EnumSet e)]
forall k a. EnumMap k a -> [a]
EnumMap.elems (EnumMap e (IntMap (EnumSet e)) -> [IntMap (EnumSet e)])
-> (EnumMap e (Set bits) -> EnumMap e (IntMap (EnumSet e)))
-> EnumMap e (Set bits)
-> [IntMap (EnumSet e)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
   (e -> Set bits -> IntMap (EnumSet e))
-> EnumMap e (Set bits) -> EnumMap e (IntMap (EnumSet e))
forall k a b. Enum k => (k -> a -> b) -> EnumMap k a -> EnumMap k b
EnumMap.mapWithKey (EnumSet e -> Set bits -> IntMap (EnumSet e)
forall bits b. C bits => b -> Set bits -> IntMap b
constIntMapFromBits (EnumSet e -> Set bits -> IntMap (EnumSet e))
-> (e -> EnumSet e) -> e -> Set bits -> IntMap (EnumSet e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. e -> EnumSet e
forall k. Enum k => k -> EnumSet k
EnumSet.singleton)


constIntMap :: b -> IntSet -> IntMap b
constIntMap :: forall b. b -> IntSet -> IntMap b
constIntMap b
a = [(Key, b)] -> IntMap b
forall a. [(Key, a)] -> IntMap a
IntMap.fromAscList ([(Key, b)] -> IntMap b)
-> (IntSet -> [(Key, b)]) -> IntSet -> IntMap b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> [Key] -> [(Key, b)]
forall b a. b -> [a] -> [(a, b)]
attach b
a ([Key] -> [(Key, b)]) -> (IntSet -> [Key]) -> IntSet -> [(Key, b)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntSet -> [Key]
IntSet.toAscList

transposeIntSet :: (Enum e) => EnumMap e IntSet -> IntMap (EnumSet e)
transposeIntSet :: forall e. Enum e => EnumMap e IntSet -> IntMap (EnumSet e)
transposeIntSet =
   (EnumSet e -> EnumSet e -> EnumSet e)
-> [IntMap (EnumSet e)] -> IntMap (EnumSet e)
forall (f :: * -> *) a.
Foldable f =>
(a -> a -> a) -> f (IntMap a) -> IntMap a
IntMap.unionsWith EnumSet e -> EnumSet e -> EnumSet e
forall k. EnumSet k -> EnumSet k -> EnumSet k
EnumSet.union ([IntMap (EnumSet e)] -> IntMap (EnumSet e))
-> (EnumMap e IntSet -> [IntMap (EnumSet e)])
-> EnumMap e IntSet
-> IntMap (EnumSet e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. EnumMap e (IntMap (EnumSet e)) -> [IntMap (EnumSet e)]
forall k a. EnumMap k a -> [a]
EnumMap.elems (EnumMap e (IntMap (EnumSet e)) -> [IntMap (EnumSet e)])
-> (EnumMap e IntSet -> EnumMap e (IntMap (EnumSet e)))
-> EnumMap e IntSet
-> [IntMap (EnumSet e)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
   (e -> IntSet -> IntMap (EnumSet e))
-> EnumMap e IntSet -> EnumMap e (IntMap (EnumSet e))
forall k a b. Enum k => (k -> a -> b) -> EnumMap k a -> EnumMap k b
EnumMap.mapWithKey (EnumSet e -> IntSet -> IntMap (EnumSet e)
forall b. b -> IntSet -> IntMap b
constIntMap (EnumSet e -> IntSet -> IntMap (EnumSet e))
-> (e -> EnumSet e) -> e -> IntSet -> IntMap (EnumSet e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. e -> EnumSet e
forall k. Enum k => k -> EnumSet k
EnumSet.singleton)