HasField with support for partial fields

HasField.hs

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ExplicitForAll #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeSynonymInstances #-}

module HasField where

import Data.Functor.Const (Const(..))
import Data.Functor.Identity (Identity(..))
import Data.Kind (Type)


-- This demonstrates a version of HasField that supports (limited) type-changing
-- update and nicely handles partial fields (giving the option to produce affine
-- traversals for them, while still yielding lenses for total fields).
--
-- The downside is that HasField has 6-7 parameters...


-- | Suggestive type synonyms. We should perhaps use new data types here.
type Total   = Const
type Partial = Either

-- | Implies that the type @p@ must be 'Total' or 'Partial', and we know which.
class KnownPartiality p where
  elimP :: (p ~ Total => r) -> (p ~ Partial => r) -> r

instance KnownPartiality Total where
  elimP x _ = x

instance KnownPartiality Partial where
  elimP _ y = y


-- | @HasField x s t a b p@ means record type @s@ has a field @x :: a@ that is
-- either total (if @p ~ Total@) or partial (if @p ~ Partial@), and setting the
-- field with a value of type @b@ produces a record type @t@.
class KnownPartiality p
    => HasField (x :: k) (s :: Type) (t :: Type) (a :: Type) (b :: Type) (p :: Type -> Type -> Type)
    | x s -> a p, x t -> b p, x s b -> t, x t a -> s where
  hasField :: s -> p (b -> t, a) t



type HasTotalField x s t a b = HasField x s t a b Total
type HasTotalField' x s a = HasField x s s a a Total


-- | Here's an example of the generated instances for a type with a total and a
-- partial field.
data T a = Foo { x :: Int, y :: a } | Bar { x :: Int }
  deriving Show

instance HasField "x" (T a) (T a) Int Int Total where
  hasField r = Const (\x' -> r{x=x'}, x r)

instance HasField "y" (T a) (T b) a b Partial where
  hasField r@(Foo{y=y}) = Left (\y' -> r{y=y'}, y)
  hasField Bar{x=x}     = Right (Bar{x=x})



-- These are just standard definitions of van Laarhoven lenses, for
-- illustration.  The same technique works with other lens representations, and
-- can generate affine traversals if the representation supports them.

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


-- The following combinators could be defined by a lens/optics library:

-- | Get the value of a total field.
getField :: forall x s t a b . HasTotalField x s t a b => s -> a
getField = snd . getConst . hasField @_ @x @s @t @a @b

-- | Get the value of a partial (or total) field, if it is present.
lookupField :: forall x s t a b p . HasField x s t a b p => s -> Maybe a
lookupField = elimP @p (Just . snd . getConst) (either (Just . snd) (const Nothing)) . hasField @_ @x @s @t @a @b

-- | Set the value of a field; has no effect if the field is absent.
setField :: forall x s t a b p . HasField x s t a b p => s -> b -> t
setField r = elimP @p (fst . getConst) (either fst const) (hasField @_ @x @s @t @a @b r)

-- | Map over the value of a field; has no effect if the field is absent.
overField :: forall x s t a b p . HasField x s t a b p => (a -> b) -> s -> t
overField f r = elimP @p (\(Const (g, a)) -> g (f a))
                         (either (\(g, a) -> g (f a)) id) (hasField @_ @x @s @t @a @b r)


-- | Get the lens corresponding to a total field.
fieldLens :: forall x s t a b . HasTotalField x s t a b => Lens s t a b
fieldLens = \ g r -> set r <$> g (get r)
  where
    get = getField @x @s @t @a @b
    set = setField @x @s @t @a @b

-- | Get the traversal corresponding to a partial field.  We could get an affine
-- traversal instead, if our lens representation supported it.
fieldTraversal' :: forall x s t a b . HasField x s t a b Partial => Traversal s t a b
fieldTraversal' = \ g r -> either (\ (_, a) -> set r <$> g a) pure (look r)
  where
    look = hasField @_ @x @s @t @a @b
    set  = setField @x @s @t @a @b

-- | Get the traversal corresponding to a field, regardless of whether it is
-- partial or total.
fieldTraversal :: forall x s t a b p . HasField x s t a b p => Traversal s t a b
fieldTraversal = elimP @p (fieldLens @x) (fieldTraversal' @x)



-- | Compute the kind of optic corresponding to a partiality: this will be
-- 'Lens' for 'Total' and 'Traversal' for 'Partial'.
type Optic p s t a b = forall f . C p f => (a -> f b) -> (s -> f t)
type Optic' p s a = Optic p s s a a

type family C p where
  C Total   = Functor
  C Partial = Applicative


-- | Get the optic corresponding to a field, producing a lens or traversal as
-- appropriate.
fieldOptic :: forall x s t a b p . HasField x s t a b p => Optic p s t a b
fieldOptic = elimP @p (fieldLens @x) (fieldTraversal @x)


eg1 :: Lens' (T a) Int
eg1 = fieldOptic @"x"

eg2 :: Traversal (T a) (T b) a b
eg2 = fieldOptic @"y"