module Data.Graph.Comfort.Map where

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 :: forall k v. (Ord k, Show k) => Caller -> Map k v -> k -> v
checkedLookup Caller
c Map k v
m k
k =
  case forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup k
k Map k v
m of
    Maybe v
Nothing ->
      forall a. HasCallStack => Caller -> a
error forall a b. (a -> b) -> a -> b
$ Caller
"checkedLookup error in " forall a. [a] -> [a] -> [a]
++ Caller
c forall a. [a] -> [a] -> [a]
++ Caller
"\n"
              forall a. [a] -> [a] -> [a]
++ Caller
"key: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> Caller
show k
k  forall a. [a] -> [a] -> [a]
++ Caller
"\n"
              forall a. [a] -> [a] -> [a]
++ Caller
"keys in map:\n" forall a. [a] -> [a] -> [a]
++ [Caller] -> Caller
unlines (forall a b. (a -> b) -> [a] -> [b]
map forall a. Show a => a -> Caller
show (forall k a. Map k a -> [k]
Map.keys Map k v
m)) forall a. [a] -> [a] -> [a]
++ Caller
"\n"
    Just v
x -> v
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 :: forall k a b c.
Ord k =>
Caller -> (a -> b -> c) -> Map k a -> Map k b -> Map k c
checkedZipWith Caller
caller a -> b -> c
f Map k a
ma Map k b
mb =
  if forall k a. Map k a -> Set k
Map.keysSet Map k a
ma forall a. Eq a => a -> a -> Bool
== forall k a. Map k a -> Set k
Map.keysSet Map k b
mb
    then forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
Map.intersectionWith a -> b -> c
f Map k a
ma Map k b
mb
    else forall a. HasCallStack => Caller -> a
error forall a b. (a -> b) -> a -> b
$
            Caller
"checkedZipWith called by function " forall a. [a] -> [a] -> [a]
++ Caller
caller forall a. [a] -> [a] -> [a]
++
            Caller
": key sets differ"


reverse :: (Ord b) => Map a b -> Map b a
reverse :: forall b a. Ord b => Map a b -> Map b a
reverse = forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> (b, a)
swap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> [(k, a)]
Map.toList

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

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

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



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

uncurry ::
   (Ord k) =>
   Caller ->
   (k0 -> k1 -> k) ->
   Map k0 (Map k1 v) -> Map k v
uncurry :: forall k k0 k1 v.
Ord k =>
Caller -> (k0 -> k1 -> k) -> Map k0 (Map k1 v) -> Map k v
uncurry Caller
caller k0 -> k1 -> k
f =
   forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
(a -> a -> a) -> f (Map k a) -> Map k a
Map.unionsWith (forall a. HasCallStack => Caller -> a
error forall a b. (a -> b) -> a -> b
$ Caller
caller forall a. [a] -> [a] -> [a]
++ Caller
".uncurry: duplicate key") forall b c a. (b -> c) -> (a -> b) -> a -> c
.
   forall k a. Map k a -> [a]
Map.elems forall b c a. (b -> c) -> (a -> b) -> a -> c
.
   forall k a b. (k -> a -> b) -> Map k a -> Map k b
Map.mapWithKey (forall k2 k1 a. Ord k2 => (k1 -> k2) -> Map k1 a -> Map k2 a
Map.mapKeys forall b c a. (b -> c) -> (a -> b) -> a -> c
. k0 -> k1 -> k
f)

flip ::
   (Ord k0, Ord k1) =>
   Map k0 (Map k1 a) -> Map k1 (Map k0 a)
flip :: forall k0 k1 a.
(Ord k0, Ord k1) =>
Map k0 (Map k1 a) -> Map k1 (Map k0 a)
flip =
   forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
(a -> a -> a) -> f (Map k a) -> Map k a
Map.unionsWith (forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith (forall a. HasCallStack => Caller -> a
error forall a b. (a -> b) -> a -> b
$ Caller
"Map.flip: duplicate key")) forall b c a. (b -> c) -> (a -> b) -> a -> c
.
   forall k a. Map k a -> [a]
Map.elems forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a b. (k -> a -> b) -> Map k a -> Map k b
Map.mapWithKey (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. k -> a -> Map k a
Map.singleton)


mapMaybeKeys ::
   (Ord k1) =>
   (k0 -> Maybe k1) ->
   Map k0 a -> Map k1 a
mapMaybeKeys :: forall k1 k0 a. Ord k1 => (k0 -> Maybe k1) -> Map k0 a -> Map k1 a
mapMaybeKeys k0 -> Maybe k1
f =
   forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (\(k0
k,a
a) -> forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a b c. (a -> b -> c) -> b -> a -> c
P.flip (,) a
a) forall a b. (a -> b) -> a -> b
$ k0 -> Maybe k1
f k0
k) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Map k a -> [(k, a)]
Map.toList


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