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

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

import           Control.Lens

import qualified Data.Strict.IntMap as IM

import           Data.Strict.IntMap (IntMap)


#if !MIN_VERSION_lens(5,0,0)
instance FunctorWithIndex Int IntMap where
  imap = IM.mapWithKey
  {-# INLINE imap #-}

instance FoldableWithIndex Int IntMap where
  ifoldMap = IM.foldMapWithKey
  {-# INLINE ifoldMap #-}
  ifoldr   = IM.foldrWithKey
  {-# INLINE ifoldr #-}
  ifoldl'  = IM.foldlWithKey' . flip
  {-# INLINE ifoldl' #-}

instance TraversableWithIndex Int IntMap where
  itraverse = IM.traverseWithKey
  {-# INLINE itraverse #-}
#endif

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

instance At (IntMap a) where
  at :: Index (IntMap a) -> Lens' (IntMap a) (Maybe (IxValue (IntMap a)))
at Index (IntMap a)
k Maybe (IxValue (IntMap a)) -> f (Maybe (IxValue (IntMap a)))
f = (Maybe a -> f (Maybe a)) -> Key -> IntMap a -> f (IntMap a)
forall (f :: * -> *) a.
Functor f =>
(Maybe a -> f (Maybe a)) -> Key -> IntMap a -> f (IntMap a)
IM.alterF Maybe a -> f (Maybe a)
Maybe (IxValue (IntMap a)) -> f (Maybe (IxValue (IntMap a)))
f Key
Index (IntMap a)
k
  {-# INLINE at #-}

instance AsEmpty (IntMap a) where
  _Empty :: p () (f ()) -> p (IntMap a) (f (IntMap a))
_Empty = IntMap a -> (IntMap a -> Bool) -> Prism' (IntMap a) ()
forall a. a -> (a -> Bool) -> Prism' a ()
nearly IntMap a
forall a. IntMap a
IM.empty IntMap a -> Bool
forall a. IntMap a -> Bool
IM.null
  {-# INLINE _Empty #-}

instance Each (IntMap a) (IntMap b) a b where
  each :: (a -> f b) -> IntMap a -> f (IntMap b)
each = (a -> f b) -> IntMap a -> f (IntMap b)
forall (f :: * -> *) a b.
Traversable f =>
IndexedTraversal Key (f a) (f b) a b
traversed
  {-# INLINE each #-}

instance (t ~ IntMap a') => Rewrapped (IntMap a) t
instance Wrapped (IntMap a) where
  type Unwrapped (IntMap a) = [(Int, a)]
  _Wrapped' :: p (Unwrapped (IntMap a)) (f (Unwrapped (IntMap a)))
-> p (IntMap a) (f (IntMap a))
_Wrapped' = (IntMap a -> [(Key, a)])
-> ([(Key, a)] -> IntMap a)
-> Iso (IntMap a) (IntMap a) [(Key, a)] [(Key, a)]
forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso IntMap a -> [(Key, a)]
forall a. IntMap a -> [(Key, a)]
IM.toAscList [(Key, a)] -> IntMap a
forall a. [(Key, a)] -> IntMap a
IM.fromList
  {-# INLINE _Wrapped' #-}

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

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