Map.hs

module Data.Graph.Comfort.Map where

import qualified Data.NonEmpty.Set as NonEmptySet
import qualified Data.Set as Set
import qualified Data.Map as Map
import Data.Set (Set)
import Data.Map (Map)
import Data.Tuple.HT (swap)
import Data.Maybe (mapMaybe)

import qualified Prelude as P
import Prelude hiding (curry, uncurry, flip)


-- | New improved ugly version with caller function name
type Caller = String


checkedLookup ::
  (Ord k, Show k) => Caller -> Map k v -> k -> v
checkedLookup c m k =
  case Map.lookup k m of
    Nothing ->
      error $ "checkedLookup error in " ++ c ++ "\n"
              ++ "key: " ++ show k  ++ "\n"
              ++ "keys in map:\n" ++ unlines (map show (Map.keys m)) ++ "\n"
    Just x -> x

{- |
The set of keys must be equal and this is checked dynamically.
-}
checkedZipWith ::
  (Ord k) =>
  Caller ->
  (a -> b -> c) ->
  Map k a -> Map k b -> Map k c
checkedZipWith caller f ma mb =
  if Map.keysSet ma == Map.keysSet mb
    then Map.intersectionWith f ma mb
    else error $
            "checkedZipWith called by function " ++ caller ++
            ": key sets differ"


reverse :: (Ord b) => Map a b -> Map b a
reverse = Map.fromList . map swap . Map.toList

-- Map.fromSet is available from containers-0.5
fromSet ::
   (Ord key) => (key -> a) -> Set key -> Map key a
fromSet f = Map.fromAscList . map (\k -> (k, f k)) . Set.toAscList

differenceSet ::
   (Ord key) => Map key a -> Set key -> Map key a
differenceSet m s = Map.difference m (fromSet (const ()) s)

intersectionSet ::
   (Ord key) => Map key a -> Set key -> Map key a
intersectionSet m s = Map.intersection m (fromSet (const ()) s)



curry ::
   (Ord k0, Ord k1) =>
   Caller ->
   (k -> (k0, k1)) ->
   Map k a -> Map k0 (Map k1 a)
curry caller f =
   Map.unionsWith (Map.unionWith (error $ caller ++ ".curry: duplicate key")) .
   Map.elems .
   Map.mapWithKey
      (\k a ->
         case f k of
            (k0, k1) -> Map.singleton k0 $ Map.singleton k1 a)

uncurry ::
   (Ord k) =>
   Caller ->
   (k0 -> k1 -> k) ->
   Map k0 (Map k1 v) -> Map k v
uncurry caller f =
   Map.unionsWith (error $ caller ++ ".uncurry: duplicate key") .
   Map.elems .
   Map.mapWithKey (Map.mapKeys . f)

flip ::
   (Ord k0, Ord k1) =>
   Map k0 (Map k1 a) -> Map k1 (Map k0 a)
flip =
   Map.unionsWith (Map.unionWith (error $ "Map.flip: duplicate key")) .
   Map.elems . Map.mapWithKey (fmap . Map.singleton)


mapMaybeKeys ::
   (Ord k1) =>
   (k0 -> Maybe k1) ->
   Map k0 a -> Map k1 a
mapMaybeKeys f =
   Map.fromList . mapMaybe (\(k,a) -> fmap (P.flip (,) a) $ f k) . Map.toList


compose :: (Ord a, Ord b) => Map b c -> Map a b -> Map a c
compose bc ab = Map.mapMaybe (P.flip Map.lookup bc) ab

transpose :: (Ord k, Ord a) => Map k a -> Map a (NonEmptySet.T k)
transpose =
   Map.foldl (Map.unionWith NonEmptySet.union) Map.empty .
   Map.mapWithKey
      (\n root -> Map.singleton root $ NonEmptySet.singleton n)