-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Multimap.Conversions
-- Maintainer  :  Ziyang Liu <free@cofree.io>
--
-- Conversions between 'Multimap' and 'SetMultimap'.
module Data.Multimap.Conversions (
  toMultimapAsc
  , toMultimapDesc
  , toSetMultimap
) where

import qualified Data.List.NonEmpty as NonEmpty
import qualified Data.Map as Map
import Data.Multimap.Internal (Multimap (..))
import Data.Multimap.Set.Internal (SetMultimap (..))
import qualified Data.Set as Set

-- | Convert a t'Data.Multimap.Set.SetMultimap' to a t'Data.Multimap.Multimap' where the values of each key
-- are in ascending order.
--
-- > toMultimapAsc (Data.Multimap.Set.fromList [(1,'a'),(1,'b'),(2,'c')]) === Data.Multimap.fromList [(1,'a'),(1,'b'),(2,'c')]
toMultimapAsc :: SetMultimap k a -> Multimap k a
toMultimapAsc :: SetMultimap k a -> Multimap k a
toMultimapAsc (SetMultimap (Map k (Set a)
m, Size
sz)) = (Map k (NonEmpty a), Size) -> Multimap k a
forall k a. (Map k (NonEmpty a), Size) -> Multimap k a
Multimap (Map k (NonEmpty a)
m', Size
sz)
  where
    m' :: Map k (NonEmpty a)
m' = (Set a -> Maybe (NonEmpty a))
-> Map k (Set a) -> Map k (NonEmpty a)
forall a b k. (a -> Maybe b) -> Map k a -> Map k b
Map.mapMaybe ([a] -> Maybe (NonEmpty a)
forall a. [a] -> Maybe (NonEmpty a)
NonEmpty.nonEmpty ([a] -> Maybe (NonEmpty a))
-> (Set a -> [a]) -> Set a -> Maybe (NonEmpty a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set a -> [a]
forall a. Set a -> [a]
Set.toAscList) Map k (Set a)
m

-- | Convert a t'Data.Multimap.Set.SetMultimap' to a t'Data.Multimap.Multimap' where the values of each key
-- are in descending order.
--
-- > toMultimapDesc (Data.Multimap.Set.fromList [(1,'a'),(1,'b'),(2,'c')]) === Data.Multimap.fromList [(1,'b'),(1,'a'),(2,'c')]
toMultimapDesc :: SetMultimap k a -> Multimap k a
toMultimapDesc :: SetMultimap k a -> Multimap k a
toMultimapDesc (SetMultimap (Map k (Set a)
m, Size
sz)) = (Map k (NonEmpty a), Size) -> Multimap k a
forall k a. (Map k (NonEmpty a), Size) -> Multimap k a
Multimap (Map k (NonEmpty a)
m', Size
sz)
  where
    m' :: Map k (NonEmpty a)
m' = (Set a -> Maybe (NonEmpty a))
-> Map k (Set a) -> Map k (NonEmpty a)
forall a b k. (a -> Maybe b) -> Map k a -> Map k b
Map.mapMaybe ([a] -> Maybe (NonEmpty a)
forall a. [a] -> Maybe (NonEmpty a)
NonEmpty.nonEmpty ([a] -> Maybe (NonEmpty a))
-> (Set a -> [a]) -> Set a -> Maybe (NonEmpty a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set a -> [a]
forall a. Set a -> [a]
Set.toDescList) Map k (Set a)
m

-- | Convert a t'Data.Multimap.Multimap' to a t'Data.Multimap.Set.SetMultimap'.
--
-- > toSetMultimap (Data.Multimap.fromList [(1,'a'),(1,'b'),(2,'c')]) === Data.Multimap.Set.fromList [(1,'a'),(1,'b'),(2,'c')]
toSetMultimap :: Ord a => Multimap k a -> SetMultimap k a
toSetMultimap :: Multimap k a -> SetMultimap k a
toSetMultimap (Multimap (Map k (NonEmpty a)
m, Size
sz)) = (Map k (Set a), Size) -> SetMultimap k a
forall k a. (Map k (Set a), Size) -> SetMultimap k a
SetMultimap (Map k (Set a)
m', Size
sz)
  where
    m' :: Map k (Set a)
m' = (NonEmpty a -> Set a) -> Map k (NonEmpty a) -> Map k (Set a)
forall a b k. (a -> b) -> Map k a -> Map k b
Map.map ([a] -> Set a
forall a. Ord a => [a] -> Set a
Set.fromList ([a] -> Set a) -> (NonEmpty a -> [a]) -> NonEmpty a -> Set a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NonEmpty a -> [a]
forall a. NonEmpty a -> [a]
NonEmpty.toList) Map k (NonEmpty a)
m