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

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

instance Each (IntMap a) (IntMap b) a b where
  each :: Traversal (IntMap a) (IntMap b) a b
each = 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' :: Iso' (IntMap a) (Unwrapped (IntMap a))
_Wrapped' = forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso forall a. IntMap a -> [(Key, a)]
IM.toAscList forall a. [(Key, a)] -> IntMap a
IM.fromList
  {-# INLINE _Wrapped' #-}

instance TraverseMin Int IntMap where
  traverseMin :: forall v. IndexedTraversal' Key (IntMap v) v
traverseMin p v (f v)
f IntMap v
m = case forall a. IntMap a -> Maybe ((Key, a), IntMap a)
IM.minViewWithKey IntMap v
m of
    Just ((Key
k,v
a), IntMap 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 forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \v
v -> forall a. (a -> Maybe a) -> IntMap a -> IntMap a
IM.updateMin (forall a b. a -> b -> a
const (forall a. a -> Maybe a
Just v
v)) IntMap v
m
    Maybe ((Key, v), IntMap v)
Nothing     -> forall (f :: * -> *) a. Applicative f => a -> f a
pure IntMap v
m
  {-# INLINE traverseMin #-}

instance TraverseMax Int IntMap where
  traverseMax :: forall v. IndexedTraversal' Key (IntMap v) v
traverseMax p v (f v)
f IntMap v
m = case forall a. IntMap a -> Maybe ((Key, a), IntMap a)
IM.maxViewWithKey IntMap v
m of
    Just ((Key
k,v
a), IntMap 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 forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \v
v -> forall a. (a -> Maybe a) -> IntMap a -> IntMap a
IM.updateMax (forall a b. a -> b -> a
const (forall a. a -> Maybe a
Just v
v)) IntMap v
m
    Maybe ((Key, v), IntMap v)
Nothing     -> forall (f :: * -> *) a. Applicative f => a -> f a
pure IntMap v
m
  {-# INLINE traverseMax #-}