Lens from Scratch

Main.hs

{-# LANGUAGE MultiParamTypeClasses, FunctionalDependencies, FlexibleInstances, Rank2Types #-}
import Control.Applicative
import Control.Monad
import Data.Functor.Identity
import Data.Foldable
import Data.Monoid
import Data.Tagged

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 = dimap id

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

instance Profunctor Tagged where
  dimap _ g = Tagged . g . unTagged

class (Profunctor p) => Choice p where
  left' :: p a b -> p (Either a c) (Either b c)
  left' = dimap (either Right Left) (either Right Left) . right'

  right' :: p a b -> p (Either c a) (Either c b)
  right' = dimap (either Right Left) (either Right Left) . left'

instance Choice (->) where
  left' k (Left a) = Left $ k a
  left' _ (Right c) = Right c

instance Choice Tagged where
  left' = Tagged . Left . unTagged

-- Equality < Iso
type Equality s t a b = forall p f. p a (f b) -> p s (f t)

simple :: Equality a a a a
simple = id

-- Iso < Lens
type Iso s t a b = forall p f. (Profunctor p, Functor f) => p a (f b) -> p s (f t)

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

enum :: Enum a => Iso Int Int a a
enum = iso toEnum fromEnum

curried :: Iso ((a, b) -> c) ((d, e) -> f) (a -> b -> c) (d -> e -> f)
curried = iso curry uncurry

reversed :: Iso String String String String
reversed = iso reverse reverse

-- Lens < Getter, Setter
type Lens s t a b = forall f. Functor f => (a -> f b) -> s -> f t

lens :: (s -> a) -> (s -> b -> t) -> Lens s t a b
lens g h = \f s -> fmap (h s) (f (g s))

-- Traversal > Lens
type Traversal s t a b = forall f. Applicative f => (a -> f b) -> s -> f t

traverseOf :: Applicative f => Traversal s t a b -> (a -> f b) -> s -> f t
traverseOf = id

both :: Traversal (a,a) (b,b) a b
both = \k (x,y) -> (,) <$> k x <*> k y

-- Prism > Lens
type Prism s t a b = forall p f. (Choice p, Applicative f) => p a (f b) -> p s (f t)

prism :: (b -> t) -> (s -> Either t a) -> Prism s t a b
prism bt sta = dimap sta (either pure (fmap bt)) . right'

_Left :: Prism (Either a c) (Either b c) a b
_Left = prism Left (either Right (Left . Right))

-- Review > Prism
type Review s a = forall p f. (Applicative f) => p a (f a) -> p s (f s)
type AReview s a = Tagged a (Identity a) -> Tagged s (Identity s)

re :: AReview s a -> Getting r a s
re r = to (runIdentity . unTagged . r . Tagged . Identity)

review :: AReview s a -> a -> s
review r a = a ^. re r

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

infixl 8 ^.., ^?
(^..) :: s -> Getting (Endo [a]) s a -> [a]
s ^.. l = (appEndo $ getConst $ l (Const . Endo . (:)) s) []

(^?) :: s -> Getting (First a) s a -> Maybe a
s ^? l = getFirst $ getConst $ l (Const . First . Just) s

folding :: Foldable f => (s -> f a) -> Fold s a
folding h = \k s -> traverse_ k (h s) *> pure s

-- Getter
type Getting r s a = (a -> Const r a) -> s -> Const r s

(^.) :: s -> Getting a s a -> a
s ^. l = getConst (l Const s)

to :: (s -> a) -> Getting r s a
to f = \k -> Const . getConst . k . f

-- Setter
type Setting s t a b = (a -> Identity b) -> s -> Identity t

infixr 4 .~, %~
(.~) :: Setting s t a b -> b -> s -> t
(.~) l = (runIdentity .) . (l . const . Identity)

(%~) :: Setting s t a b -> (a -> b) -> s -> t
(%~) l f = runIdentity . l (Identity . f)

sets :: ((a -> b) -> s -> t) -> Setting s t a b
sets h = \k -> Identity . h (runIdentity . k)

-- Tuple
class TupleIndex s t a b | s -> a, t -> b, s b -> t, t a -> s where
  _1 :: Lens s t a b

instance TupleIndex (a,b) (a',b) a a' where
  _1 = lens (\(a,_) -> a) (\(_,y) b -> (b,y))
instance TupleIndex (a,b,c) (a',b,c) a a' where
  _1 = lens (\(a,_,_) -> a) (\(_,y,z) b -> (b,y,z))

-- Each
class Each s t a b | s -> a, t -> b, s b -> t, t a -> s where
  each :: Traversal s t a b

instance Each [a] [b] a b where
  each = traverse

-- Cons
class Cons s t a b | s -> a, t -> b, s b -> t, t a -> s where
  _Cons :: Prism s t (a,s) (b,t)

instance Cons [a] [b] a b where
  _Cons = prism (uncurry (:)) $ \ass -> case ass of
    (a:as) -> Right (a,as)
    [] -> Left []

infixr 5 <|
(<|) :: Cons s s a a => a -> s -> s
a <| s = review _Cons (a,s)

_head :: Cons s s a a => Traversal s s a a
_head = \k s -> _Cons (\(a,s') -> (,) <$> k a <*> pure s') s

-- Ix
ix :: Int -> Lens [a] [a] a a
ix n = lens (!! n) (\ts x -> take n ts ++ [x] ++ drop (n+1) ts)

main = do
  putStrLn "Lens from Scratch"
  print $ ("Hello","World") ^. _1
  print $ [1..10] ^. ix 7
  print $ _1 .~ "Hello" $ (1,())
  print $ ix 5 .~ 0 $ [1..10]
  traverseOf each print [1,2,3]
  print $ both %~ (*10) $ (1,2)
  print $ ("hello","world") ^. both
  print $ ("hello","world") ^.. both
  print $ [[1,2],[3]] ^.. traverse . traverse
  print $ [1..10] ^.. folding tail
  print $ (97 ^. enum :: Char)
  print $ reversed %~ ('d' :) $ "live"
  print $ Left "Hello" ^. _Left
  print $ (review _Left "hogehoge" :: Either String ())
  print $ 0 <| [1,2,3]
  print $ [1,2,3] ^? _head
  print $ _head .~ 3 $ [0,1]