{-# 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 k -> Map k a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
Index (Map k a)
k Map k a
m of
     Just a
v  -> IxValue (Map k a) -> f (IxValue (Map k a))
f a
IxValue (Map k a)
v f a -> (a -> Map k a) -> f (Map k a)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \a
v' -> k -> a -> Map k a -> Map k a
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert k
Index (Map k a)
k a
v' Map k a
m
     Maybe a
Nothing -> Map k a -> f (Map k a)
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 = (Maybe v -> f (Maybe v)) -> k -> Map k v -> f (Map k v)
forall (f :: * -> *) k a.
(Functor f, Ord k) =>
(Maybe a -> f (Maybe a)) -> k -> Map k a -> f (Map k a)
M.alterF Maybe v -> f (Maybe v)
Maybe (IxValue (Map k v)) -> f (Maybe (IxValue (Map k v)))
f k
Index (Map k v)
k
  {-# INLINE at #-}

instance AsEmpty (Map k a) where
  _Empty :: p () (f ()) -> p (Map k a) (f (Map k a))
_Empty = Map k a -> (Map k a -> Bool) -> Prism' (Map k a) ()
forall a. a -> (a -> Bool) -> Prism' a ()
nearly Map k a
forall k a. Map k a
M.empty Map k a -> Bool
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 :: (a -> f b) -> Map c a -> f (Map d b)
each = (a -> f b) -> Map c a -> f (Map d b)
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' :: p (Unwrapped (Map k a)) (f (Unwrapped (Map k a)))
-> p (Map k a) (f (Map k a))
_Wrapped' = (Map k a -> [(k, a)])
-> ([(k, a)] -> Map k a)
-> Iso (Map k a) (Map k a) [(k, a)] [(k, a)]
forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso Map k a -> [(k, a)]
forall k a. Map k a -> [(k, a)]
M.toAscList [(k, a)] -> Map k a
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList
  {-# INLINE _Wrapped' #-}

instance Ord k => TraverseMin k (Map k) where
  traverseMin :: p v (f v) -> Map k v -> f (Map k v)
traverseMin p v (f v)
f Map k v
m = case Map k v -> Maybe ((k, v), Map k v)
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
_) -> p v (f v) -> k -> v -> f 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 f v -> (v -> Map k v) -> f (Map k v)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \v
v -> (v -> Maybe v) -> Map k v -> Map k v
forall a k. (a -> Maybe a) -> Map k a -> Map k a
M.updateMin (Maybe v -> v -> Maybe v
forall a b. a -> b -> a
const (v -> Maybe v
forall a. a -> Maybe a
Just v
v)) Map k v
m
    Maybe ((k, v), Map k v)
Nothing          -> Map k v -> f (Map k v)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Map k v
m
  {-# INLINE traverseMin #-}

instance Ord k => TraverseMax k (Map k) where
  traverseMax :: p v (f v) -> Map k v -> f (Map k v)
traverseMax p v (f v)
f Map k v
m = case Map k v -> Maybe ((k, v), Map k v)
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
_) -> p v (f v) -> k -> v -> f 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 f v -> (v -> Map k v) -> f (Map k v)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \v
v -> (v -> Maybe v) -> Map k v -> Map k v
forall a k. (a -> Maybe a) -> Map k a -> Map k a
M.updateMax (Maybe v -> v -> Maybe v
forall a b. a -> b -> a
const (v -> Maybe v
forall a. a -> Maybe a
Just v
v)) Map k v
m
    Maybe ((k, v), Map k v)
Nothing          -> Map k v -> f (Map k v)
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 :: IndexedGetting i (Map i a) s a -> s -> Map i a
toMapOf IndexedGetting i (Map i a) s a
l = IndexedGetting i (Map i a) s a
-> (i -> a -> Map i a) -> s -> Map i a
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 i -> a -> Map i a
forall k a. k -> a -> Map k a
M.singleton