module Data.Key (
Key
, Keyed(..)
, (<#$>)
, keyed
, apWithKey
, Index(..)
, (!)
, Lookup(..)
, lookupDefault
, FoldableWithKey(..)
, foldrWithKey'
, foldlWithKey'
, foldrWithKeyM
, foldlWithKeyM
, traverseWithKey_
, forWithKey_
, mapWithKeyM_
, forWithKeyM_
, concatMapWithKey
, anyWithKey
, allWithKey
, findWithKey
, FoldableWithKey1(..)
, traverseWithKey1_
, forWithKey1_
, foldMapWithKeyDefault1
, TraversableWithKey(..)
, forWithKey
, forWithKeyM
, mapAccumWithKeyL
, mapAccumWithKeyR
, mapWithKeyDefault
, foldMapWithKeyDefault
, TraversableWithKey1(..)
, foldMapWithKey1Default
, module Data.Foldable
, module Data.Traversable
, module Data.Semigroup
, module Data.Semigroup.Foldable
, module Data.Semigroup.Traversable
) where
import Control.Applicative
import Control.Comonad.Trans.Traced
import Control.Monad.Trans.Identity
import Control.Monad.Trans.Reader
import qualified Data.Array as Array
import Data.Array (Array)
import Data.Functor.Identity
import Data.Functor.Bind
import Data.Functor.Compose
import Data.Foldable
import Data.IntMap (IntMap)
import qualified Data.IntMap as IntMap
import Data.Ix hiding (index)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe (listToMaybe)
import Data.Monoid as Monoid
import Data.Semigroup
import Data.Semigroup.Foldable
import Data.Semigroup.Traversable
import Data.Sequence (Seq, ViewL(..), viewl)
import qualified Data.Sequence as Seq
import Data.Traversable
import Prelude hiding (lookup)
type family Key (f :: * -> *)
class Functor f => Keyed f where
mapWithKey :: (Key f -> a -> b) -> f a -> f b
infixl 4 <#$>
(<#$>) :: Keyed f => (Key f -> a -> b) -> f a -> f b
(<#$>) = mapWithKey
keyed :: Keyed f => f a -> f (Key f, a)
keyed = mapWithKey (,)
apWithKey :: (Keyed f, Applicative f) => f (Key f -> a -> b) -> f a -> f b
apWithKey ff fa = mapWithKey (\k f -> f k) ff <*> fa
class Index f where
index :: f a -> Key f -> a
(!) :: Index f => f a -> Key f -> a
(!) = index
class Lookup f where
lookup :: Key f -> f a -> Maybe a
lookupDefault :: Index f => Key f -> f a -> Maybe a
lookupDefault k t = Just (index t k)
class Foldable t => FoldableWithKey t where
toIndexedList :: t a -> [(Key t, a)]
toIndexedList = foldrWithKey (\k v t -> (k,v):t) []
foldMapWithKey :: Monoid m => (Key t -> a -> m) -> t a -> m
foldMapWithKey f = foldrWithKey (\k v -> mappend (f k v)) mempty
foldrWithKey :: (Key t -> a -> b -> b) -> b -> t a -> b
foldrWithKey f z t = appEndo (foldMapWithKey (\k v -> Endo (f k v)) t) z
foldlWithKey :: (b -> Key t -> a -> b) -> b -> t a -> b
foldlWithKey f z t = appEndo (getDual (foldMapWithKey (\k a -> Dual (Endo (\b -> f b k a))) t)) z
foldrWithKey' :: FoldableWithKey t => (Key t -> a -> b -> b) -> b -> t a -> b
foldrWithKey' f z0 xs = foldlWithKey f' id xs z0
where f' k key x z = k $! f key x z
foldlWithKey' :: FoldableWithKey t => (b -> Key t -> a -> b) -> b -> t a -> b
foldlWithKey' f z0 xs = foldrWithKey f' id xs z0
where f' key x k z = k $! f z key x
foldrWithKeyM :: (FoldableWithKey t, Monad m) => (Key t -> a -> b -> m b) -> b -> t a -> m b
foldrWithKeyM f z0 xs = foldlWithKey f' return xs z0
where f' k key x z = f key x z >>= k
foldlWithKeyM :: (FoldableWithKey t, Monad m) => (b -> Key t -> a -> m b) -> b -> t a -> m b
foldlWithKeyM f z0 xs = foldrWithKey f' return xs z0
where f' key x k z = f z key x >>= k
traverseWithKey_ :: (FoldableWithKey t, Applicative f) => (Key t -> a -> f b) -> t a -> f ()
traverseWithKey_ f = foldrWithKey (fmap (*>) . f) (pure ())
forWithKey_ :: (FoldableWithKey t, Applicative f) => t a -> (Key t -> a -> f b) -> f ()
forWithKey_ = flip traverseWithKey_
mapWithKeyM_ :: (FoldableWithKey t, Monad m) => (Key t -> a -> m b) -> t a -> m ()
mapWithKeyM_ f = foldrWithKey (fmap (>>) . f) (return ())
forWithKeyM_ :: (FoldableWithKey t, Monad m) => t a -> (Key t -> a -> m b) -> m ()
forWithKeyM_ = flip mapWithKeyM_
concatMapWithKey :: FoldableWithKey t => (Key t -> a -> [b]) -> t a -> [b]
concatMapWithKey = foldMapWithKey
anyWithKey :: FoldableWithKey t => (Key t -> a -> Bool) -> t a -> Bool
anyWithKey p = getAny . foldMapWithKey (fmap Any . p)
allWithKey :: FoldableWithKey t => (Key t -> a -> Bool) -> t a -> Bool
allWithKey p = getAll . foldMapWithKey (fmap All . p)
findWithKey :: FoldableWithKey t => (Key t -> a -> Bool) -> t a -> Maybe a
findWithKey p = Monoid.getFirst . foldMapWithKey (\k x -> Monoid.First (if p k x then Just x else Nothing) )
class (Foldable1 t, FoldableWithKey t) => FoldableWithKey1 t where
foldMapWithKey1 :: Semigroup m => (Key t -> a -> m) -> t a -> m
newtype Act f a = Act { getAct :: f a }
instance Apply f => Semigroup (Act f a) where
Act a <> Act b = Act (a .> b)
instance Functor f => Functor (Act f) where
fmap f (Act a) = Act (f <$> a)
b <$ Act a = Act (b <$ a)
traverseWithKey1_ :: (FoldableWithKey1 t, Apply f) => (Key t -> a -> f b) -> t a -> f ()
traverseWithKey1_ f = (<$) () . getAct . foldMapWithKey1 (fmap Act . f)
forWithKey1_ :: (FoldableWithKey1 t, Apply f) => t a -> (Key t -> a -> f b) -> f ()
forWithKey1_ = flip traverseWithKey1_
foldMapWithKeyDefault1 :: (FoldableWithKey1 t, Monoid m) => (Key t -> a -> m) -> t a -> m
foldMapWithKeyDefault1 f = unwrapMonoid . foldMapWithKey (fmap WrapMonoid . f)
class (Keyed t, FoldableWithKey t, Traversable t) => TraversableWithKey t where
traverseWithKey :: Applicative f => (Key t -> a -> f b) -> t a -> f (t b)
mapWithKeyM :: Monad m => (Key t -> a -> m b) -> t a -> m (t b)
mapWithKeyM f = unwrapMonad . traverseWithKey (fmap WrapMonad . f)
forWithKey :: (TraversableWithKey t, Applicative f) => t a -> (Key t -> a -> f b) -> f (t b)
forWithKey = flip traverseWithKey
forWithKeyM :: (TraversableWithKey t, Monad m) => t a -> (Key t -> a -> m b) -> m (t b)
forWithKeyM = flip mapWithKeyM
newtype StateL s a = StateL { runStateL :: s -> (s, a) }
instance Functor (StateL s) where
fmap f (StateL k) = StateL $ \ s ->
let (s', v) = k s in (s', f v)
instance Applicative (StateL s) where
pure x = StateL (\ s -> (s, x))
StateL kf <*> StateL kv = StateL $ \ s ->
let (s', f) = kf s
(s'', v) = kv s'
in (s'', f v)
mapAccumWithKeyL :: TraversableWithKey t => (Key t -> a -> b -> (a, c)) -> a -> t b -> (a, t c)
mapAccumWithKeyL f s t = runStateL (traverseWithKey (\k b -> StateL (\a -> f k a b)) t) s
newtype StateR s a = StateR { runStateR :: s -> (s, a) }
instance Functor (StateR s) where
fmap f (StateR k) = StateR $ \ s ->
let (s', v) = k s in (s', f v)
instance Applicative (StateR s) where
pure x = StateR (\ s -> (s, x))
StateR kf <*> StateR kv = StateR $ \ s ->
let (s', v) = kv s
(s'', f) = kf s'
in (s'', f v)
mapAccumWithKeyR :: TraversableWithKey t => (Key t -> a -> b -> (a, c)) -> a -> t b -> (a, t c)
mapAccumWithKeyR f s t = runStateR (traverseWithKey (\k b -> StateR (\a -> f k a b)) t) s
mapWithKeyDefault :: TraversableWithKey t => (Key t -> a -> b) -> t a -> t b
mapWithKeyDefault f = runIdentity . traverseWithKey (fmap Identity . f)
foldMapWithKeyDefault :: (TraversableWithKey t, Monoid m) => (Key t -> a -> m) -> t a -> m
foldMapWithKeyDefault f = getConst . traverseWithKey (fmap Const . f)
class (Traversable1 t, FoldableWithKey1 t, TraversableWithKey t) => TraversableWithKey1 t where
traverseWithKey1 :: Apply f => (Key t -> a -> f b) -> t a -> f (t b)
foldMapWithKey1Default :: (TraversableWithKey1 t, Semigroup m) => (Key t -> a -> m) -> t a -> m
foldMapWithKey1Default f = getConst . traverseWithKey1 (\k -> Const . f k)
type instance Key Identity = ()
instance Index Identity where
index (Identity a) _ = a
instance Lookup Identity where
lookup _ (Identity a) = Just a
instance Keyed Identity where
mapWithKey f = Identity . f () . runIdentity
instance FoldableWithKey Identity where
foldrWithKey f z (Identity a) = f () a z
instance FoldableWithKey1 Identity where
foldMapWithKey1 f (Identity a) = f () a
instance TraversableWithKey Identity where
traverseWithKey f (Identity a) = Identity <$> f () a
instance TraversableWithKey1 Identity where
traverseWithKey1 f (Identity a) = Identity <$> f () a
type instance Key (IdentityT m) = Key m
instance Index m => Index (IdentityT m) where
index (IdentityT m) i = index m i
instance Lookup m => Lookup (IdentityT m) where
lookup i (IdentityT m) = lookup i m
instance Keyed m => Keyed (IdentityT m) where
mapWithKey f = IdentityT . mapWithKey f . runIdentityT
instance FoldableWithKey m => FoldableWithKey (IdentityT m) where
foldrWithKey f z (IdentityT m) = foldrWithKey f z m
instance FoldableWithKey1 m => FoldableWithKey1 (IdentityT m) where
foldMapWithKey1 f (IdentityT m) = foldMapWithKey1 f m
instance TraversableWithKey m => TraversableWithKey (IdentityT m) where
traverseWithKey f (IdentityT a) = IdentityT <$> traverseWithKey f a
instance TraversableWithKey1 m => TraversableWithKey1 (IdentityT m) where
traverseWithKey1 f (IdentityT a) = IdentityT <$> traverseWithKey1 f a
type instance Key ((->)a) = a
instance Keyed ((->)a) where
mapWithKey = (<*>)
instance Index ((->)a) where
index = id
instance Lookup ((->)a) where
lookup i f = Just (f i)
type instance Key (ReaderT e m) = (e, Key m)
instance Keyed m => Keyed (ReaderT e m) where
mapWithKey f (ReaderT m) = ReaderT $ \k -> mapWithKey (f . (,) k) (m k)
instance Index m => Index (ReaderT e m) where
index (ReaderT f) (e,k) = index (f e) k
instance Lookup m => Lookup (ReaderT e m) where
lookup (e,k) (ReaderT f) = lookup k (f e)
type instance Key (TracedT s w) = (s, Key w)
instance Keyed w => Keyed (TracedT s w) where
mapWithKey f = TracedT . mapWithKey (\k' g k -> f (k, k') (g k)) . runTracedT
instance Index w => Index (TracedT s w) where
index (TracedT w) (e,k) = index w k e
instance Lookup w => Lookup (TracedT s w) where
lookup (e,k) (TracedT w) = ($ e) <$> lookup k w
type instance Key IntMap = Int
instance Keyed IntMap where
mapWithKey = IntMap.mapWithKey
instance FoldableWithKey IntMap where
foldrWithKey = IntMap.foldWithKey
instance TraversableWithKey IntMap where
traverseWithKey f = fmap IntMap.fromDistinctAscList . traverse (\(k, v) -> (,) k <$> f k v) . IntMap.toAscList
instance Index IntMap where
index = (IntMap.!)
instance Lookup IntMap where
lookup = IntMap.lookup
type instance Key (Compose f g) = (Key f, Key g)
instance (Keyed f, Keyed g) => Keyed (Compose f g) where
mapWithKey f = Compose . mapWithKey (\k -> mapWithKey (f . (,) k)) . getCompose
instance (Index f, Index g) => Index (Compose f g) where
index (Compose fg) (i,j) = index (index fg i) j
instance (Lookup f, Lookup g) => Lookup (Compose f g) where
lookup (i,j) (Compose fg) = lookup i fg >>= lookup j
instance (FoldableWithKey f, FoldableWithKey m) => FoldableWithKey (Compose f m) where
foldMapWithKey f = foldMapWithKey (\k -> foldMapWithKey (f . (,) k)) . getCompose
instance (FoldableWithKey1 f, FoldableWithKey1 m) => FoldableWithKey1 (Compose f m) where
foldMapWithKey1 f = foldMapWithKey1 (\k -> foldMapWithKey1 (f . (,) k)) . getCompose
instance (TraversableWithKey f, TraversableWithKey m) => TraversableWithKey (Compose f m) where
traverseWithKey f = fmap Compose . traverseWithKey (\k -> traverseWithKey (f . (,) k)) . getCompose
instance (TraversableWithKey1 f, TraversableWithKey1 m) => TraversableWithKey1 (Compose f m) where
traverseWithKey1 f = fmap Compose . traverseWithKey1 (\k -> traverseWithKey1 (f . (,) k)) . getCompose
type instance Key [] = Int
instance Keyed [] where
mapWithKey f0 xs0 = go f0 xs0 0 where
go _ [] _ = []
go f (x:xs) n = f n x : (go f xs $! n)
instance FoldableWithKey [] where
foldrWithKey f0 z0 xs0 = go f0 z0 xs0 0 where
go _ z [] _ = z
go f z (x:xs) n = f n x (go f z xs $! n)
instance TraversableWithKey [] where
traverseWithKey f0 xs0 = go f0 xs0 0 where
go _ [] _ = pure []
go f (x:xs) n = (:) <$> f n x <*> (go f xs $! (n + 1))
instance Index [] where
index = (!!)
instance Lookup [] where
lookup = fmap listToMaybe . drop
type instance Key Seq = Int
instance Index Seq where
index = Seq.index
instance Lookup Seq where
lookup i s = case viewl (Seq.drop i s) of
EmptyL -> Nothing
a :< _ -> Just a
instance Keyed Seq where
mapWithKey = Seq.mapWithIndex
instance FoldableWithKey Seq where
foldrWithKey = Seq.foldrWithIndex
instance TraversableWithKey Seq where
traverseWithKey f = fmap Seq.fromList . traverseWithKey f . toList
type instance Key (Map k) = k
instance Keyed (Map k) where
mapWithKey = Map.mapWithKey
instance Ord k => Index (Map k) where
index = (Map.!)
instance Ord k => Lookup (Map k) where
lookup = Map.lookup
instance FoldableWithKey (Map k) where
foldrWithKey = Map.foldrWithKey
instance TraversableWithKey (Map k) where
traverseWithKey f = fmap Map.fromDistinctAscList . traverse (\(k, v) -> (,) k <$> f k v) . Map.toAscList
type instance Key (Array i) = i
instance Ix i => Keyed (Array i) where
mapWithKey f arr = Array.listArray (Array.bounds arr) $ map (uncurry f) $ Array.assocs arr
instance Ix i => Index (Array i) where
index = (Array.!)
instance Ix i => Lookup (Array i) where
lookup i arr
| inRange (Array.bounds arr) i = Just (arr Array.! i)
| otherwise = Nothing
instance Ix i => FoldableWithKey (Array i) where
foldrWithKey f z = Prelude.foldr (uncurry f) z . Array.assocs
instance Ix i => TraversableWithKey (Array i) where
traverseWithKey f arr = Array.listArray (Array.bounds arr) <$> traverse (uncurry f) (Array.assocs arr)