my personal lens library

Lens.hs

{-# LANGUAGE RankNTypes, BlockArguments #-}

module Lens
  ( Lens
  , MonoLens
  , view
  , (...)
  , over
  , set
  , lens
  , use
  , ($=)
  , (#=)
  , (+=)
  , (-=)
  , (^=)
  , (|=)
  , (&=)
  , ($~)
  , (#~)
  , (+~)
  , (-~)
  , (^~)
  , (|~)
  , (&~)
  , Indexable(..)
  , at
  , (<#=)
  , (<+=)
  , (<-=)
  , (<^=)
  , (<|=)
  , (<&=)
  ) where

import Data.Bits
import Control.Monad.State
import Data.Functor.Const
import Data.Functor.Identity
import qualified Data.Sequence as Seq

type Lens s t a b = forall f. Functor f => (a -> f b) -> (s -> f t)
type MonoLens s a = Lens s s a a

view :: Lens s t a b -> s -> a
view l = getConst . l Const

{-# INLINE view #-}

infixl 8 ...

(...) :: s -> Lens s t a b -> a
(...) v l = view l v

{-# INLINE (...) #-}

over :: Lens s t a b -> (a -> b) -> s -> t
over l f = runIdentity . l (Identity . f)

{-# INLINE over #-}

set :: Lens s t a b -> b -> s -> t
set l a = runIdentity . l (Identity . const a)

{-# INLINE set #-}

lens :: (s -> a) -> (s -> b -> t) -> Lens s t a b
lens sa sbt afb s = sbt s <$> afb (sa s)

{-# INLINE lens #-}

use :: MonadState s m => Lens s s a a -> m a
use l = gets $ getConst . l Const

{-# INLINE use #-}

infix 4 $=, #=, +=, -=, ^=, |=, &=
infix 4 $~, #~, +~, -~, ^~, |~, &~
infixr 2 <#=, <+=, <-=, <^=, <|=, <&=

($=) :: MonadState s m => Lens s s a b -> (a -> b) -> m ()
($=) l f = modify $ over l f

{-# INLINE ($=) #-}

(#=) :: MonadState s m => Lens s s a b -> b -> m ()
(#=) l v = modify $ over l $ const v

{-# INLINE (#=) #-}

(+=) :: (Num a, MonadState s m) => Lens s s a a -> a -> m ()
(+=) l v = modify $ over l (v +)

{-# INLINE (+=) #-}

(-=) :: (Num a, MonadState s m) => Lens s s a a -> a -> m ()
(-=) l v = modify $ over l (v -)

{-# INLINE (-=) #-}

(^=) :: (Bits a, MonadState s m) => Lens s s a a -> a -> m ()
(^=) l v = modify $ over l (xor v)

{-# INLINE (^=) #-}

(|=) :: (Bits a, MonadState s m) => Lens s s a a -> a -> m ()
(|=) l v = modify $ over l (v .|.)

{-# INLINE (|=) #-}

(&=) :: (Bits a, MonadState s m) => Lens s s a a -> a -> m ()
(&=) l v = modify $ over l (v .&.)

{-# INLINE (&=) #-}

($~) :: Lens s s a b -> (a -> b) -> s -> s
($~) l f = over l f

{-# INLINE ($~) #-}

(#~) :: Lens s s a b -> b -> s -> s
(#~) l v = over l $ const v

{-# INLINE (#~) #-}

(+~) :: Num a => Lens s s a a -> a -> s -> s
(+~) l v = over l (v +)

{-# INLINE (+~) #-}

(-~) :: Num a => Lens s s a a -> a -> s -> s
(-~) l v = over l (v -)

{-# INLINE (-~) #-}

(^~) :: Bits a => Lens s s a a -> a -> s -> s
(^~) l v = over l (xor v)

{-# INLINE (^~) #-}

(|~) :: Bits a => Lens s s a a -> a -> s -> s
(|~) l v = over l (v .|.)

{-# INLINE (|~) #-}

(&~) :: Bits a => Lens s s a a -> a -> s -> s
(&~) l v = over l (v .&.)

{-# INLINE (&~) #-}

(<#=) :: MonadState s m => Lens s s a b -> m b -> m ()
(<#=) l v = v >>= (l #=)

{-# INLINE (<#=) #-}

(<+=) :: (Num a, MonadState s m) => Lens s s a a -> m a -> m ()
(<+=) l v = modify =<< over l <$> fmap (+) v

{-# INLINE (<+=) #-}

(<-=) :: (Num a, MonadState s m) => Lens s s a a -> m a -> m ()
(<-=) l v = modify =<< over l <$> fmap (-) v

{-# INLINE (<-=) #-}

(<^=) :: (Bits a, MonadState s m) => Lens s s a a -> m a -> m ()
(<^=) l v = modify =<< over l <$> fmap xor v

{-# INLINE (<^=) #-}

(<|=) :: (Bits a, MonadState s m) => Lens s s a a -> m a -> m ()
(<|=) l v = modify =<< over l <$> fmap (.|.) v

{-# INLINE (<|=) #-}

(<&=) :: (Bits a, MonadState s m) => Lens s s a a -> m a -> m ()
(<&=) l v = modify =<< over l <$> fmap (.&.) v

{-# INLINE (<&=) #-}

class Indexable c where
  getAt :: Int -> c a -> a
  putAt :: Int -> c a -> a -> c a

instance Indexable Seq.Seq where
  getAt idx val = Seq.index val idx
  {-# INLINE getAt #-}
  putAt idx old new = Seq.update idx new old
  {-# INLINE putAt #-}

at :: Indexable c => Int -> MonoLens (c a) a
at idx = lens (getAt idx) (putAt idx)

{-# INLINE at #-}