Optics.hs

{-# LANGUAGE FlexibleInstances, RankNTypes #-}

module Optics where

import Control.Category (Category, (>>>))
import qualified Control.Category as Cat
import Control.Monad ((<=<))
import Control.Monad.Identity (Identity(Identity), runIdentity)
import Data.Function ((&))
import Data.Functor.Const (Const(Const), getConst)

-- Lens

data Lens s a = Lens
    { lGet :: s -> a
    , lSet :: a -> s -> s
    }

instance Category Lens where
    id = Lens id const
    l1 . l2 = Lens
        { lGet = lGet l1 . lGet l2
        , lSet = \x r -> lSet l2 (lSet l1 x (lGet l2 r)) r
        }

class FromLens o where
    fromLens :: Lens s a -> o s a

instance FromLens Lens where
    fromLens = id

-- Getter

newtype Getter s a = Getter (s -> a)

instance Category Getter where
    id = Getter id
    Getter f . Getter g = Getter (f . g)

class FromLens o => FromGetter o where
    fromGetter :: Getter s a -> o s a

instance FromGetter Getter where
    fromGetter = id

instance FromLens Getter where
    fromLens = Getter . lGet

(^.) :: s -> Getter s a -> a
s ^. Getter f = f s
infixl 8 ^.

to :: FromGetter o => (s -> a) -> o s a
to = fromGetter . Getter

-- Traversal

newtype Traversal s a = Traversal (forall f. Applicative f => (a -> f a) -> s -> f s)

instance Category Traversal where
    id = Traversal id
    Traversal f . Traversal g = Traversal (g . f)

class FromLens o => FromTraversal o where
    fromTraversal :: Traversal s a -> o s a

instance FromTraversal Traversal where
    fromTraversal = id

instance FromLens Traversal where
    fromLens l = Traversal (\f s -> (\x -> lSet l x s) <$> f (lGet l s))

-- Fold

newtype Fold s a = Fold (s -> [a])

instance Category Fold where
    id = Fold pure
    Fold f . Fold g = Fold (f <=< g)

class (FromGetter o, FromTraversal o) => FromFold o where
    fromFold :: Fold s a -> o s a

instance FromFold Fold where
    fromFold = id

instance FromGetter Fold where
    fromGetter (Getter g) = Fold (pure . g)

instance FromTraversal Fold where
    fromTraversal (Traversal t) = Fold (getConst . t (Const . pure))

instance FromLens Fold where
    fromLens = fromGetter . fromLens

(^..) :: s -> Fold s a -> [a]
s ^.. Fold f = f s
infixl 8 ^..

-- Setter

newtype Setter s a = Setter ((a -> a) -> s -> s)

instance Category Setter where
    id = Setter id
    Setter f . Setter g = Setter (g . f)

class FromTraversal o => FromSetter o where
    fromSetter :: Setter s a -> o s a

instance FromSetter Setter where
    fromSetter = id

instance FromTraversal Setter where
    fromTraversal (Traversal t) = Setter (\f -> runIdentity . t (Identity . f))

instance FromLens Setter where
    fromLens l = Setter (\f s -> lSet l (f (lGet l s)) s)

(%~) :: Setter s a -> (a -> a) -> s -> s
Setter f %~ g = f g
infixr 4 %~

(.~) :: Setter s a -> a -> s -> s
s .~ x = s %~ const x
infixr 4 .~

sets :: FromSetter o => ((a -> a) -> s -> s) -> o s a
sets = fromSetter . Setter

-- Implementations

_1 :: FromLens o => o (a, b) a
_1 = fromLens $ Lens
    { lGet = fst
    , lSet = \x (_, y) -> (x, y)
    }

_2 :: FromLens o => o (a, b) b
_2 = fromLens $ Lens
    { lGet = snd
    , lSet = \y (x, _) -> (x, y)
    }

_Just :: FromTraversal o => o (Maybe a) a
_Just = fromTraversal $ Traversal traverse

_Left :: FromTraversal o => o (Either a b) a
_Left = fromTraversal $ Traversal traverseLeft
  where
    traverseLeft f (Left l) = Left <$> f l
    traverseLeft f s@(Right _) = pure s

_Right :: FromTraversal o => o (Either a b) b
_Right = fromTraversal $ Traversal traverse

-- Examples

five :: Int
five = (5, 1) ^. _1

left6 :: Either Int Bool
left6 = Left 5 & _Left %~ (+ 1)

seven :: Int
seven = (6, (7, 8)) ^. (_2 >>> _1)

one8 :: [Int]
one8 = (1, Just 8) ^.. (_2 >>> _Just)