Optics.hs

{-# LANGUAGE RankNTypes
           , BlockArguments
           , GADTs
           , MultiParamTypeClasses
           , OverloadedLabels
           , DataKinds
           , PolyKinds
           , FunctionalDependencies
           , AllowAmbiguousTypes
           , TypeApplications
           , ScopedTypeVariables
           , FlexibleInstances
           , TypeOperators
           , LiberalTypeSynonyms
           , ConstraintKinds
           , TupleSections
           #-}

module Optics where

import Data.List
import Data.Tuple
import Control.Monad.State
import Data.Functor.Identity
import Data.Functor.Const

-- Helper Functions

swapEither :: Either a b -> Either b a
swapEither = either Right Left

-- Interfaces

class Profunctor p where
  dimap :: (a -> b) -> (c -> d) -> p b c -> p a d
  dimap f g = lmap f . rmap g

  lmap :: (a -> b) -> p b c -> p a c
  lmap f = dimap f id

  rmap :: (b -> c) -> p a b -> p a c
  rmap f = dimap id f

  {-# MINIMAL dimap | (lmap , rmap) #-}

class Profunctor p => Strong p where
  first :: p a b -> p (a, c) (b, c)
  first = dimap swap swap . second

  second :: p a b -> p (c, a) (c, b)
  second = dimap swap swap . first

  {-# MINIMAL first | second #-}

class Profunctor p => Choice p where
  left :: p a b -> p (Either a c) (Either b c)
  left = dimap swapEither swapEither . right

  right :: p a b -> p (Either c a) (Either c b)
  right = dimap swapEither swapEither . left

  {-# MINIMAL left | right #-}

class Profunctor p => Monoidal p where
  par :: p a b -> p c d -> p (a, c) (b, d)
  empty :: p () ()

-- Wrapper Types

data Star f a b
  = Star { runStar :: a -> f b }

data Tagged a b
  = Tagged { unTagged :: b }

data Exchange a b s t
  = Exchange (s -> a) (b -> t)

-- (->) Implementations

instance Profunctor (->) where
  dimap f g p = g . p . f

instance Strong (->) where
  first f (a, c) = (f a, c)
  second f (a, c) = (a, f c)

instance Choice (->) where
  right f = either Left (Right . f)
  left f = either (Left . f) Right

instance Monoidal (->) where
  par f g = \(a, c) -> (f a, g c)
  empty = id

-- Star Implementations

instance Functor f => Profunctor (Star f) where
  dimap f g (Star h) = Star (fmap g . h . f)

instance Functor f => Strong (Star f) where
  first (Star f) = Star \(a, c) -> fmap (, c) (f a)
  second (Star f) = Star \(c, b) -> fmap (c ,) (f b)

instance Applicative f => Choice (Star f) where
  left (Star f) = Star (either (fmap Left . f) (fmap Right . pure))
  right (Star f) = Star (either (fmap Left . pure) (fmap Right . f))

instance Applicative f => Monoidal (Star f) where
  par (Star f) (Star g) = Star \(a, b) -> (,) <$> f a <*> g b
  empty = Star pure

-- Tagged Implementations

instance Profunctor Tagged where
  dimap _ g (Tagged a) = Tagged (g a)

instance Choice Tagged where
  left (Tagged a) = Tagged (Left a)
  right (Tagged a) = Tagged (Right a)

instance Monoidal Tagged where
  par (Tagged a) (Tagged b) = Tagged (a, b)
  empty = Tagged ()

-- Exchange Implementations

instance Profunctor (Exchange a b) where
  dimap f g (Exchange sa bt) = Exchange (sa . f) (g . bt)

instance Functor (Exchange a b s) where
  fmap f (Exchange sa bt) = Exchange sa (f . bt)

-- Helper Types

type Affine f = (Choice f, Strong f)
type Simple f s a = f s s a a

-- Optic

type Optic p s t a b = p a b -> p s t

-- Iso, and Lens

type Iso s t a b = forall p. Profunctor p => Optic p s t a b
type Lens s t a b = forall p. Strong p => Optic p s t a b

-- Prism, RawTraversal and Traversal

type Prism s t a b = forall p. Choice p => Optic p s t a b
type RawTraversal s t a b = forall p. Affine p => Optic p s t a b
type Traversal s t a b = forall p. Affine p => Optic p s t a b

-- Getter and Setter

type Getter s a = Optic (Star (Const a)) s s a a
type Setter s t a b = Optic (Star Identity) s t a b

-- Constructors

iso :: (s -> a) -> (b -> t) -> Iso s t a b
iso sa bt = dimap sa bt

lens :: (s -> a) -> (s -> b -> t) -> Lens s t a b
lens getter setter pab = dimap get set (first pab)
  where get s = (getter s, s)
        set (b, s) = setter s b

prism :: (b -> t) -> (s -> Either t a) -> Prism s t a b
prism setter getter = dimap getter (either id setter) . right

rawTraversal :: (s -> Either t a) -> (s -> b -> t) -> RawTraversal s t a b
rawTraversal getter setter pab = dimap get set (first (right pab))
  where get s = (getter s, s)
        set (bt, s) = either id (setter s) bt

-- Combinators

view :: Getter s a -> s -> a
view l = getConst . runStar (l (Star Const))

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

set :: Setter s t a b -> b -> s -> t
set l v = runIdentity . runStar (l (Star (Identity . const v)))

infixl 8 ^.

(^.) :: s -> Getter s a -> a
(^.) = flip view

infix 4 %~, .~

(%~) :: Setter s t a b -> (a -> b) -> s -> t
(%~) l f = over l f

(.~) :: Setter s t a b -> b -> s -> t
(.~) l v = over l (const v)

infix 4 %=, .=

(%=) :: MonadState s m => Setter s s a b -> (a -> b) -> m ()
(%=) l f = modify (l %~ f)

(.=) :: MonadState s m => Setter s s a b -> b -> m ()
(.=) l v = modify (l .~ v)

-- Optics

_fst :: Lens (a, c) (b, c) a b
_fst = lens fst \(_, b) a -> (a, b)

_snd :: Lens (c, a) (c, b) a b
_snd = lens snd \(a, _) b -> (a, b)

_head :: Simple Lens [a] a
_head = lens head \(_ : xs) x -> x : xs

_tail :: Simple Lens [a] [a]
_tail = lens tail \(x : _) xs -> x : xs

swapped :: Simple Iso (a, b) (b, a)
swapped = iso swap swap

flipped :: Simple Iso (a -> b -> c) (b -> a -> c)
flipped = iso flip flip

reversed :: Simple Iso [a] [a]
reversed = iso reverse reverse

simple :: Simple Iso a a
simple = id