{-# LANGUAGE CPP                   #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE UndecidableInstances  #-}

{-# OPTIONS_GHC -fno-warn-orphans #-}
module Data.Strict.Map.Lens
  ( toMapOf
  ) where

import           Control.Lens

import qualified Data.Strict.Map as M

import           Data.Strict.Map (Map)


#if !MIN_VERSION_lens(5,0,0)
instance FunctorWithIndex k (Map k) where
  imap = M.mapWithKey
  {-# INLINE imap #-}

instance FoldableWithIndex k (Map k) where
  ifoldMap = M.foldMapWithKey
  {-# INLINE ifoldMap #-}
  ifoldr   = M.foldrWithKey
  {-# INLINE ifoldr #-}
  ifoldl'  = M.foldlWithKey' . flip
  {-# INLINE ifoldl' #-}

instance TraversableWithIndex k (Map k) where
  itraverse = M.traverseWithKey
  {-# INLINE itraverse #-}
#endif

type instance Index (Map k v) = k
type instance IxValue (Map k v) = v

instance Ord k => Ixed (Map k a) where
  ix :: Index (Map k a) -> Traversal' (Map k a) (IxValue (Map k a))
ix Index (Map k a)
k IxValue (Map k a) -> f (IxValue (Map k a))
f Map k a
m = case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup Index (Map k a)
k Map k a
m of
     Just a
v  -> IxValue (Map k a) -> f (IxValue (Map k a))
f a
v forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \a
v' -> forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert Index (Map k a)
k a
v' Map k a
m
     Maybe a
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure Map k a
m
  {-# INLINE ix #-}

instance Ord k => At (Map k v) where
  at :: Index (Map k v) -> Lens' (Map k v) (Maybe (IxValue (Map k v)))
at Index (Map k v)
k Maybe (IxValue (Map k v)) -> f (Maybe (IxValue (Map k v)))
f = forall (f :: * -> *) k a.
(Functor f, Ord k) =>
(Maybe a -> f (Maybe a)) -> k -> Map k a -> f (Map k a)
M.alterF Maybe (IxValue (Map k v)) -> f (Maybe (IxValue (Map k v)))
f Index (Map k v)
k
  {-# INLINE at #-}

instance AsEmpty (Map k a) where
  _Empty :: Prism' (Map k a) ()
_Empty = forall a. a -> (a -> Bool) -> Prism' a ()
nearly forall k a. Map k a
M.empty forall k a. Map k a -> Bool
M.null
  {-# INLINE _Empty #-}

instance (c ~ d) => Each (Map c a) (Map d b) a b where
  each :: Traversal (Map c a) (Map d b) a b
each = forall (f :: * -> *) a b.
Traversable f =>
IndexedTraversal Int (f a) (f b) a b
traversed
  {-# INLINE each #-}

instance (t ~ Map k' a', Ord k) => Rewrapped (Map k a) t
instance Ord k => Wrapped (Map k a) where
  type Unwrapped (Map k a) = [(k, a)]
  _Wrapped' :: Iso' (Map k a) (Unwrapped (Map k a))
_Wrapped' = forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso forall k a. Map k a -> [(k, a)]
M.toAscList forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
  {-# INLINE _Wrapped' #-}

instance Ord k => TraverseMin k (Map k) where
  traverseMin :: forall v. IndexedTraversal' k (Map k v) v
traverseMin p v (f v)
f Map k v
m = case forall k a. Map k a -> Maybe ((k, a), Map k a)
M.minViewWithKey Map k v
m of
    Just ((k
k, v
a), Map k v
_) -> forall i (p :: * -> * -> *) a b.
Indexable i p =>
p a b -> i -> a -> b
indexed p v (f v)
f k
k v
a forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \v
v -> forall a k. (a -> Maybe a) -> Map k a -> Map k a
M.updateMin (forall a b. a -> b -> a
const (forall a. a -> Maybe a
Just v
v)) Map k v
m
    Maybe ((k, v), Map k v)
Nothing          -> forall (f :: * -> *) a. Applicative f => a -> f a
pure Map k v
m
  {-# INLINE traverseMin #-}

instance Ord k => TraverseMax k (Map k) where
  traverseMax :: forall v. IndexedTraversal' k (Map k v) v
traverseMax p v (f v)
f Map k v
m = case forall k a. Map k a -> Maybe ((k, a), Map k a)
M.maxViewWithKey Map k v
m of
    Just ((k
k, v
a), Map k v
_) -> forall i (p :: * -> * -> *) a b.
Indexable i p =>
p a b -> i -> a -> b
indexed p v (f v)
f k
k v
a forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \v
v -> forall a k. (a -> Maybe a) -> Map k a -> Map k a
M.updateMax (forall a b. a -> b -> a
const (forall a. a -> Maybe a
Just v
v)) Map k v
m
    Maybe ((k, v), Map k v)
Nothing          -> forall (f :: * -> *) a. Applicative f => a -> f a
pure Map k v
m
  {-# INLINE traverseMax #-}

-- | Analogous to 'Data.Map.Lens.toMapOf'.
toMapOf :: IndexedGetting i (Map i a) s a -> s -> Map i a
toMapOf :: forall i a s. IndexedGetting i (Map i a) s a -> s -> Map i a
toMapOf IndexedGetting i (Map i a) s a
l = forall s (m :: * -> *) i r a.
MonadReader s m =>
IndexedGetting i r s a -> (i -> a -> r) -> m r
iviews IndexedGetting i (Map i a) s a
l forall k a. k -> a -> Map k a
M.singleton