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


reverse :: (Ord b) => Map a b -> Map b a
reverse :: Map a b -> Map b a
reverse = [(b, a)] -> Map b a
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(b, a)] -> Map b a)
-> (Map a b -> [(b, a)]) -> Map a b -> Map b a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((a, b) -> (b, a)) -> [(a, b)] -> [(b, a)]
forall a b. (a -> b) -> [a] -> [b]
map (a, b) -> (b, a)
forall a b. (a, b) -> (b, a)
swap ([(a, b)] -> [(b, a)])
-> (Map a b -> [(a, b)]) -> Map a b -> [(b, a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map a b -> [(a, b)]
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 :: (key -> a) -> Set key -> Map key a
fromSet key -> a
f = [(key, a)] -> Map key a
forall k a. Eq k => [(k, a)] -> Map k a
Map.fromAscList ([(key, a)] -> Map key a)
-> (Set key -> [(key, a)]) -> Set key -> Map key a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (key -> (key, a)) -> [key] -> [(key, a)]
forall a b. (a -> b) -> [a] -> [b]
map (\key
k -> (key
k, key -> a
f key
k)) ([key] -> [(key, a)])
-> (Set key -> [key]) -> Set key -> [(key, a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Set key -> [key]
forall a. Set a -> [a]
Set.toAscList

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

intersectionSet ::
   (Ord key) => Map key a -> Set key -> Map key a
intersectionSet :: Map key a -> Set key -> Map key a
intersectionSet Map key a
m Set key
s = Map key a -> Map key () -> Map key a
forall k a b. Ord k => Map k a -> Map k b -> Map k a
Map.intersection Map key a
m ((key -> ()) -> Set key -> Map key ()
forall key a. Ord key => (key -> a) -> Set key -> Map key a
fromSet (() -> key -> ()
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 :: Caller -> (k -> (k0, k1)) -> Map k a -> Map k0 (Map k1 a)
curry Caller
caller k -> (k0, k1)
f =
   (Map k1 a -> Map k1 a -> Map k1 a)
-> [Map k0 (Map k1 a)] -> Map k0 (Map k1 a)
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
(a -> a -> a) -> f (Map k a) -> Map k a
Map.unionsWith ((a -> a -> a) -> Map k1 a -> Map k1 a -> Map k1 a
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith (Caller -> a -> a -> a
forall a. HasCallStack => Caller -> a
error (Caller -> a -> a -> a) -> Caller -> a -> a -> a
forall a b. (a -> b) -> a -> b
$ Caller
caller Caller -> Caller -> Caller
forall a. [a] -> [a] -> [a]
++ Caller
".curry: duplicate key")) ([Map k0 (Map k1 a)] -> Map k0 (Map k1 a))
-> (Map k a -> [Map k0 (Map k1 a)]) -> Map k a -> Map k0 (Map k1 a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
   Map k (Map k0 (Map k1 a)) -> [Map k0 (Map k1 a)]
forall k a. Map k a -> [a]
Map.elems (Map k (Map k0 (Map k1 a)) -> [Map k0 (Map k1 a)])
-> (Map k a -> Map k (Map k0 (Map k1 a)))
-> Map k a
-> [Map k0 (Map k1 a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
   (k -> a -> Map k0 (Map k1 a))
-> Map k a -> Map k (Map k0 (Map k1 a))
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) -> k0 -> Map k1 a -> Map k0 (Map k1 a)
forall k a. k -> a -> Map k a
Map.singleton k0
k0 (Map k1 a -> Map k0 (Map k1 a)) -> Map k1 a -> Map k0 (Map k1 a)
forall a b. (a -> b) -> a -> b
$ k1 -> a -> Map k1 a
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 :: Caller -> (k0 -> k1 -> k) -> Map k0 (Map k1 v) -> Map k v
uncurry Caller
caller k0 -> k1 -> k
f =
   (v -> v -> v) -> [Map k v] -> Map k v
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
(a -> a -> a) -> f (Map k a) -> Map k a
Map.unionsWith (Caller -> v -> v -> v
forall a. HasCallStack => Caller -> a
error (Caller -> v -> v -> v) -> Caller -> v -> v -> v
forall a b. (a -> b) -> a -> b
$ Caller
caller Caller -> Caller -> Caller
forall a. [a] -> [a] -> [a]
++ Caller
".uncurry: duplicate key") ([Map k v] -> Map k v)
-> (Map k0 (Map k1 v) -> [Map k v]) -> Map k0 (Map k1 v) -> Map k v
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
   Map k0 (Map k v) -> [Map k v]
forall k a. Map k a -> [a]
Map.elems (Map k0 (Map k v) -> [Map k v])
-> (Map k0 (Map k1 v) -> Map k0 (Map k v))
-> Map k0 (Map k1 v)
-> [Map k v]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
   (k0 -> Map k1 v -> Map k v)
-> Map k0 (Map k1 v) -> Map k0 (Map k v)
forall k a b. (k -> a -> b) -> Map k a -> Map k b
Map.mapWithKey ((k1 -> k) -> Map k1 v -> Map k v
forall k2 k1 a. Ord k2 => (k1 -> k2) -> Map k1 a -> Map k2 a
Map.mapKeys ((k1 -> k) -> Map k1 v -> Map k v)
-> (k0 -> k1 -> k) -> k0 -> Map k1 v -> Map k v
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 :: Map k0 (Map k1 a) -> Map k1 (Map k0 a)
flip =
   (Map k0 a -> Map k0 a -> Map k0 a)
-> [Map k1 (Map k0 a)] -> Map k1 (Map k0 a)
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
(a -> a -> a) -> f (Map k a) -> Map k a
Map.unionsWith ((a -> a -> a) -> Map k0 a -> Map k0 a -> Map k0 a
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith (Caller -> a -> a -> a
forall a. HasCallStack => Caller -> a
error (Caller -> a -> a -> a) -> Caller -> a -> a -> a
forall a b. (a -> b) -> a -> b
$ Caller
"Map.flip: duplicate key")) ([Map k1 (Map k0 a)] -> Map k1 (Map k0 a))
-> (Map k0 (Map k1 a) -> [Map k1 (Map k0 a)])
-> Map k0 (Map k1 a)
-> Map k1 (Map k0 a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
   Map k0 (Map k1 (Map k0 a)) -> [Map k1 (Map k0 a)]
forall k a. Map k a -> [a]
Map.elems (Map k0 (Map k1 (Map k0 a)) -> [Map k1 (Map k0 a)])
-> (Map k0 (Map k1 a) -> Map k0 (Map k1 (Map k0 a)))
-> Map k0 (Map k1 a)
-> [Map k1 (Map k0 a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (k0 -> Map k1 a -> Map k1 (Map k0 a))
-> Map k0 (Map k1 a) -> Map k0 (Map k1 (Map k0 a))
forall k a b. (k -> a -> b) -> Map k a -> Map k b
Map.mapWithKey ((a -> Map k0 a) -> Map k1 a -> Map k1 (Map k0 a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a -> Map k0 a) -> Map k1 a -> Map k1 (Map k0 a))
-> (k0 -> a -> Map k0 a) -> k0 -> Map k1 a -> Map k1 (Map k0 a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. k0 -> a -> Map k0 a
forall k a. k -> a -> Map k a
Map.singleton)


mapMaybeKeys ::
   (Ord k1) =>
   (k0 -> Maybe k1) ->
   Map k0 a -> Map k1 a
mapMaybeKeys :: (k0 -> Maybe k1) -> Map k0 a -> Map k1 a
mapMaybeKeys k0 -> Maybe k1
f =
   [(k1, a)] -> Map k1 a
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(k1, a)] -> Map k1 a)
-> (Map k0 a -> [(k1, a)]) -> Map k0 a -> Map k1 a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((k0, a) -> Maybe (k1, a)) -> [(k0, a)] -> [(k1, a)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (\(k0
k,a
a) -> (k1 -> (k1, a)) -> Maybe k1 -> Maybe (k1, a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((k1 -> a -> (k1, a)) -> a -> k1 -> (k1, a)
forall a b c. (a -> b -> c) -> b -> a -> c
P.flip (,) a
a) (Maybe k1 -> Maybe (k1, a)) -> Maybe k1 -> Maybe (k1, a)
forall a b. (a -> b) -> a -> b
$ k0 -> Maybe k1
f k0
k) ([(k0, a)] -> [(k1, a)])
-> (Map k0 a -> [(k0, a)]) -> Map k0 a -> [(k1, a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map k0 a -> [(k0, a)]
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 :: Map b c -> Map a b -> Map a c
compose Map b c
bc Map a b
ab = (b -> Maybe c) -> Map a b -> Map a c
forall a b k. (a -> Maybe b) -> Map k a -> Map k b
Map.mapMaybe ((b -> Map b c -> Maybe c) -> Map b c -> b -> Maybe c
forall a b c. (a -> b -> c) -> b -> a -> c
P.flip b -> Map b c -> Maybe c
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Map b c
bc) Map a b
ab