-- | Data-shuffling
module Data.Cfg.Collect
  ( collectOnFirst
  , collectOnSecond
  , collectOnFirst'
  , collectOnSecond'
  ) where

import Data.List (nub)
import qualified Data.Set as S

-- | Collects a list of pairs on the first element.
collectOnFirst :: Eq a => [(a, b)] -> [(a, [b])]
collectOnFirst :: [(a, b)] -> [(a, [b])]
collectOnFirst [(a, b)]
pairs = [(a
a, a -> [b]
bsFor a
a) | a
a <- [a]
as]
  where
    as :: [a]
as = [a] -> [a]
forall a. Eq a => [a] -> [a]
nub ([a] -> [a]) -> [a] -> [a]
forall a b. (a -> b) -> a -> b
$ ((a, b) -> a) -> [(a, b)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a, b) -> a
forall a b. (a, b) -> a
fst [(a, b)]
pairs
    bsFor :: a -> [b]
bsFor a
a = [b
b | (a
a', b
b) <- [(a, b)]
pairs, a
a a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
a']

-- | Collects a list of pairs on the second element.
collectOnSecond :: Eq b => [(a, b)] -> [([a], b)]
collectOnSecond :: [(a, b)] -> [([a], b)]
collectOnSecond [(a, b)]
pairs = [(b -> [a]
asFor b
b, b
b) | b
b <- [b]
bs]
  where
    bs :: [b]
bs = [b] -> [b]
forall a. Eq a => [a] -> [a]
nub ([b] -> [b]) -> [b] -> [b]
forall a b. (a -> b) -> a -> b
$ ((a, b) -> b) -> [(a, b)] -> [b]
forall a b. (a -> b) -> [a] -> [b]
map (a, b) -> b
forall a b. (a, b) -> b
snd [(a, b)]
pairs
    asFor :: b -> [a]
asFor b
b = [a
a | (a
a, b
b') <- [(a, b)]
pairs, b
b b -> b -> Bool
forall a. Eq a => a -> a -> Bool
== b
b']

-- | Collects a list of pairs on the first element.
collectOnFirst' :: (Eq a, Ord b) => [(a, b)] -> [(a, S.Set b)]
collectOnFirst' :: [(a, b)] -> [(a, Set b)]
collectOnFirst' [(a, b)]
pairs = [(a
a, a -> Set b
bsFor a
a) | a
a <- [a]
as]
  where
    as :: [a]
as = [a] -> [a]
forall a. Eq a => [a] -> [a]
nub ([a] -> [a]) -> [a] -> [a]
forall a b. (a -> b) -> a -> b
$ ((a, b) -> a) -> [(a, b)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (a, b) -> a
forall a b. (a, b) -> a
fst [(a, b)]
pairs
    bsFor :: a -> Set b
bsFor a
a = [b] -> Set b
forall a. Ord a => [a] -> Set a
S.fromList [b
b | (a
a', b
b) <- [(a, b)]
pairs, a
a a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
a']

-- | Collects a list of pairs on the second element.
collectOnSecond' :: (Ord a, Eq b) => [(a, b)] -> [(S.Set a, b)]
collectOnSecond' :: [(a, b)] -> [(Set a, b)]
collectOnSecond' [(a, b)]
pairs = [(b -> Set a
asFor b
b, b
b) | b
b <- [b]
bs]
  where
    bs :: [b]
bs = [b] -> [b]
forall a. Eq a => [a] -> [a]
nub ([b] -> [b]) -> [b] -> [b]
forall a b. (a -> b) -> a -> b
$ ((a, b) -> b) -> [(a, b)] -> [b]
forall a b. (a -> b) -> [a] -> [b]
map (a, b) -> b
forall a b. (a, b) -> b
snd [(a, b)]
pairs
    asFor :: b -> Set a
asFor b
b = [a] -> Set a
forall a. Ord a => [a] -> Set a
S.fromList [a
a | (a
a, b
b') <- [(a, b)]
pairs, b
b b -> b -> Bool
forall a. Eq a => a -> a -> Bool
== b
b']