CoRecord.hs

-- |Module containing the sum formulation companion to 'Composite.Record's product formulation. Values of type @'CoRec' f rs@ represent a single value
-- @f r@ for one of the @r@s in @rs@. Heavily based on the great work by Anthony Cowley in Frames.
{-# LANGUAGE UndecidableInstances #-} -- for FoldRec
module Composite.CoRecord where

import Prelude
import Composite.Record (AllHave, HasInstances, (:->)(getVal, Val), reifyDicts, reifyVal, val, zipRecsWith)
import Control.Lens (Prism', Wrapped, Unwrapped, prism', review, view, _Wrapped')
import Control.Monad.Except (ExceptT, throwError, withExceptT)
import Data.Functor.Contravariant (Contravariant(contramap))
import Data.Functor.Identity (Identity(Identity), runIdentity)
import Data.Kind (Constraint, Type)
import Data.Maybe (fromMaybe)
import Data.Profunctor (dimap)
import Data.Proxy (Proxy(Proxy))
import Data.Text (Text, pack)
import Data.Vinyl.Core (Dict(Dict), Rec((:&), RNil), RMap, RecApplicative, RecordToList, ReifyConstraint, recordToList, reifyConstraint, rmap, rpure)
import Data.Vinyl.Functor (Compose(Compose, getCompose), Const(Const), (:.))
import Data.Vinyl.Lens (RElem, type (∈), type (⊆), rget, rput, rreplace)
import Data.Vinyl.TypeLevel (RecAll, RIndex)
import GHC.TypeLits (KnownSymbol, symbolVal)

-- FIXME? replace with int-index/union or at least lift ideas from there. This encoding is awkward to work with and not compositional.

-- |@CoRef f rs@ represents a single value of type @f r@ for some @r@ in @rs@.
data CoRec :: (u -> Type) -> [u] -> Type where
  -- |Witness that @r@ is an element of @rs@ using '∈' ('RElem' with 'RIndex') from Vinyl.
  CoVal :: r ∈ rs => !(f r) -> CoRec f rs

instance forall rs. (AllHave '[Show] rs, RecApplicative rs) => Show (CoRec Identity rs) where
  show (CoVal (Identity x)) = "(CoVal " ++ show' x ++ ")"
    where
      shower :: Rec (Op String) rs
      shower = reifyDicts (Proxy @'[Show]) (\ _ -> Op show)
      show' = runOp (rget shower)

instance forall rs. (RMap rs, RecAll Maybe rs Eq, RecApplicative rs, RecordToList rs, ReifyConstraint Eq Maybe rs) => Eq (CoRec Identity rs) where
  crA == crB = and . recordToList $ zipRecsWith f (toRec crA) (fieldToRec crB)
    where
      f :: forall a. (Dict Eq :. Maybe) a -> Maybe a -> Const Bool a
      f (Compose (Dict a)) b = Const $ a == b
      toRec = reifyConstraint . fieldToRec

-- |The common case of a 'CoRec' with @f ~ 'Identity'@, i.e. a regular value.
type Field = CoRec Identity

-- |Inject a value @f r@ into a @'CoRec' f rs@ given that @r@ is one of the valid @rs@.
--
-- Equivalent to 'CoVal' the constructor, but exists to parallel 'field'.
coRec :: r ∈ rs => f r -> CoRec f rs
coRec = CoVal

-- |Produce a prism for the given alternative of a 'CoRec'.
coRecPrism :: (RecApplicative rs, r ∈ rs) => Prism' (CoRec f rs) (f r)
coRecPrism = prism' CoVal (getCompose . rget . coRecToRec)

-- |Inject a value @r@ into a @'Field' rs@ given that @r@ is one of the valid @rs@.
--
-- Equivalent to @'CoVal' . 'Identity'@.
field :: r ∈ rs => r -> Field rs
field = CoVal . Identity

-- |Inject a value @a@ into a @'Field' rs@ given that @s :-> a@ is one of the valid @rs@.
--
-- Equivalent to @'CoVal' . 'Identity' . 'Val'@.
fieldVal :: forall s a rs proxy. s :-> a ∈ rs => proxy (s :-> a) -> a -> Field rs
fieldVal _ = CoVal . val @s

-- |Produce a prism for the given alternative of a 'Field'.
fieldPrism :: (RecApplicative rs, r ∈ rs) => Prism' (Field rs) r
fieldPrism = coRecPrism . dimap runIdentity (fmap Identity)

-- |Produce a prism for the given @:->@ alternative of a 'Field', given a proxy to identify which @s :-> a@ you meant.
fieldValPrism :: (RecApplicative rs, s :-> a ∈ rs) => proxy (s :-> a) -> Prism' (Field rs) a
fieldValPrism proxy = coRecPrism . dimap (getVal . reifyVal proxy . runIdentity) (fmap (Identity . Val))

-- |Apply an extraction to whatever @f r@ is contained in the given 'CoRec'.
--
-- For example @foldCoVal getConst :: CoRec (Const a) rs -> a@.
foldCoVal :: (forall (r :: u). RElem r rs (RIndex r rs) => f r -> b) -> CoRec f rs -> b
foldCoVal f (CoVal x) = f x
{-# INLINE foldCoVal #-}

-- |Map a @'CoRec' f@ to a @'CoRec' g@ using a natural transform from @f@ to @g@ (@forall x. f x -> g x@).
mapCoRec :: (forall x. f x -> g x) -> CoRec f rs -> CoRec g rs
mapCoRec f (CoVal x) = CoVal (f x)
{-# INLINE mapCoRec #-}

-- |Apply some kleisli on @h@ to the @f x@ contained in a @'CoRec' f@ and yank the @h@ outside. Like 'traverse' but for 'CoRec'.
traverseCoRec :: Functor h => (forall x. f x -> h (g x)) -> CoRec f rs -> h (CoRec g rs)
traverseCoRec f (CoVal x) = CoVal <$> f x
{-# INLINE traverseCoRec #-}

-- |Project a @'CoRec' f@ into a @'Rec' ('Maybe' ':.' f)@ where only the single @r@ held by the 'CoRec' is 'Just' in the resulting record, and all other
-- fields are 'Nothing'.
coRecToRec :: RecApplicative rs => CoRec f rs -> Rec (Maybe :. f) rs
coRecToRec (CoVal a) = rput (Compose (Just a)) (rpure (Compose Nothing))
{-# INLINE coRecToRec #-}

-- |Project a 'Field' into a @'Rec' 'Maybe'@ where only the single @r@ held by the 'Field' is 'Just' in the resulting record, and all other
-- fields are 'Nothing'.
fieldToRec :: (RMap rs, RecApplicative rs) => Field rs -> Rec Maybe rs
fieldToRec = rmap (fmap runIdentity . getCompose) . coRecToRec
{-# INLINE fieldToRec #-}

-- |Typeclass which allows folding ala 'foldMap' over a 'Rec', using a 'CoRec' as the accumulator.
class FoldRec ss ts where
  -- |Given some combining function, an initial value, and a record, visit each field of the record using the combining function to accumulate the
  -- initial value or previous accumulation with the field of the record.
  foldRec
    :: (CoRec f ss -> CoRec f ss -> CoRec f ss)
    -> CoRec f ss
    -> Rec f ts
    -> CoRec f ss

instance FoldRec ss '[] where
  foldRec _ z _ = z
  {-# INLINE foldRec #-}

instance (t ∈ ss, FoldRec ss ts) => FoldRec ss (t ': ts) where
  foldRec f z (x :& xs) = foldRec f (z `f` CoVal x) xs
  {-# INLINE foldRec #-}

-- |'foldRec' for records with at least one field that doesn't require an initial value.
foldRec1
  :: FoldRec (r ': rs) rs
  => (CoRec f (r ': rs) -> CoRec f (r ': rs) -> CoRec f (r ': rs))
  -> Rec f (r ': rs)
  -> CoRec f (r ': rs)
foldRec1 f (x :& xs) = foldRec f (CoVal x) xs
{-# INLINE foldRec1 #-}

-- |Given a @'Rec' ('Maybe' ':.' f) rs@, yield a @Just coRec@ for the first field which is @Just@, or @Nothing@ if there are no @Just@ fields in the record.
firstCoRec :: FoldRec rs rs => Rec (Maybe :. f) rs -> Maybe (CoRec f rs)
firstCoRec RNil       = Nothing
firstCoRec v@(x :& _) = traverseCoRec getCompose $ foldRec f (CoVal x) v
  where
    f c@(CoVal (Compose (Just _))) _ = c
    f _                            c = c
{-# INLINE firstCoRec #-}

-- |Given a @'Rec' 'Maybe' rs@, yield a @Just field@ for the first field which is @Just@, or @Nothing@ if there are no @Just@ fields in the record.
firstField :: (FoldRec rs rs, RMap rs) => Rec Maybe rs -> Maybe (Field rs)
firstField = firstCoRec . rmap (Compose . fmap Identity)
{-# INLINE firstField #-}

-- |Given a @'Rec' ('Maybe' ':.' f) rs@, yield a @Just coRec@ for the last field which is @Just@, or @Nothing@ if there are no @Just@ fields in the record.
lastCoRec :: FoldRec rs rs => Rec (Maybe :. f) rs -> Maybe (CoRec f rs)
lastCoRec RNil       = Nothing
lastCoRec v@(x :& _) = traverseCoRec getCompose $ foldRec f (CoVal x) v
  where
    f _ c@(CoVal (Compose (Just _))) = c
    f c                            _ = c
{-# INLINE lastCoRec #-}

-- |Given a @'Rec' 'Maybe' rs@, yield a @Just field@ for the last field which is @Just@, or @Nothing@ if there are no @Just@ fields in the record.
lastField :: (RMap rs, FoldRec rs rs) => Rec Maybe rs -> Maybe (Field rs)
lastField = lastCoRec . rmap (Compose . fmap Identity)
{-# INLINE lastField #-}

-- |Helper newtype containing a function @a -> b@ but with the type parameters flipped so @Op b@ has a consistent codomain for a varying domain.
newtype Op b a = Op { runOp :: a -> b }

-- |Given a list of constraints @cs@ required to apply some function, apply the function to whatever value @r@ (not @f r@) which the 'CoRec' contains.
onCoRec
  :: forall (cs :: [Type -> Constraint]) (f :: Type -> Type) (rs :: [Type]) (b :: Type) (proxy :: [Type -> Constraint] -> Type).
     (AllHave cs rs, Functor f, RecApplicative rs)
  => proxy cs
  -> (forall (a :: Type). HasInstances a cs => a -> b)
  -> CoRec f rs
  -> f b
onCoRec p f (CoVal x) = go <$> x
  where
    go = runOp $ rget (reifyDicts p (\ _ -> Op f) :: Rec (Op b) rs)
{-# INLINE onCoRec #-}

-- |Given a list of constraints @cs@ required to apply some function, apply the function to whatever value @r@ which the 'Field' contains.
onField
  :: forall (cs :: [Type -> Constraint]) (rs :: [Type]) (b :: Type) (proxy :: [Type -> Constraint] -> Type).
     (AllHave cs rs, RecApplicative rs)
  => proxy cs
  -> (forall (a :: Type). HasInstances a cs => a -> b)
  -> Field rs
  -> b
onField p f x = runIdentity (onCoRec p f x)
{-# INLINE onField #-}

-- |Given some target type @r@ that's a possible value of @'Field' rs@, yield @Just@ if that is indeed the value being stored by the 'Field', or @Nothing@ if
-- not.
asA :: (r ∈ rs, RMap rs, RecApplicative rs) => Field rs -> Maybe r
asA = rget . fieldToRec
{-# INLINE asA #-}

-- |An extractor function @f a -> b@ which can be passed to 'foldCoRec' to eliminate one possible alternative of a 'CoRec'.
newtype Case' f b a = Case' { unCase' :: f a -> b }

instance Functor f => Contravariant (Case' f b) where
  contramap f (Case' r) = Case' $ \x -> r (fmap f x) 

-- |A record of @Case'@ eliminators for each @r@ in @rs@ representing the pieces of a total function from @'CoRec' f@ to @b@.
type Cases' f rs b = Rec (Case' f b) rs

-- |Fold a @'CoRec' f@ using @Cases'@ which eliminate each possible value held by the 'CoRec', yielding the @b@ produced by whichever case matches.
foldCoRec :: RecApplicative (r ': rs) => Cases' f (r ': rs) b -> CoRec f (r ': rs) -> b
foldCoRec hs = go hs . coRecToRec
  where
    go :: Cases' f rs b -> Rec (Maybe :. f) rs -> b
    go (Case' f :&  _) (Compose (Just x) :& _) = f x
    go (Case' _ :& fs) (Compose Nothing  :& t) = go fs t
    go RNil            RNil                    = error "foldCoRec should be provably total, isn't"
    {-# INLINE go #-}
{-# INLINE foldCoRec #-}

-- |Fold a @'CoRec' f@ using @Cases'@ which eliminate each possible value held by the 'CoRec', yielding the @b@ produced by whichever case matches.
--
-- Equivalent to 'foldCoRec' but with its arguments flipped so it can be written @matchCoRec coRec $ cases@.
matchCoRec :: RecApplicative (r ': rs) => CoRec f (r ': rs) -> Cases' f (r ': rs) b -> b
matchCoRec = flip foldCoRec
{-# INLINE matchCoRec #-}

newtype Case b a = Case { unCase :: a -> b }

instance Contravariant (Case b) where
  contramap f (Case r) = Case $ r . f

type Cases rs b = Rec (Case b) rs

-- |Fold a 'Field' using 'Cases' which eliminate each possible value held by the 'Field', yielding the @b@ produced by whichever case matches.
foldField :: (RMap rs, RecApplicative (r ': rs)) => Cases (r ': rs) b -> Field (r ': rs) -> b
foldField hs = foldCoRec (rmap (Case' . (. runIdentity) . unCase) hs)
{-# INLINE foldField #-}

-- |Fold a 'Field' using 'Cases' which eliminate each possible value held by the 'Field', yielding the @b@ produced by whichever case matches.
--
-- Equivalent to 'foldCoRec' but with its arguments flipped so it can be written @matchCoRec coRec $ cases@.
matchField :: (RMap rs, RecApplicative (r ': rs)) => Field (r ': rs) -> Cases (r ': rs) b -> b
matchField = flip foldField
{-# INLINE matchField #-}

-- |Widen a @'CoRec' f rs@ to a @'CoRec' f ss@ given that @rs ⊆ ss@.
widenCoRec :: (FoldRec ss ss, RecApplicative rs, RecApplicative ss, rs ⊆ ss) => CoRec f rs -> CoRec f ss
widenCoRec r =
  fromMaybe (error "widenCoRec should be provably total, isn't") $
    firstCoRec (rreplace (coRecToRec r) (rpure $ Compose Nothing))

-- |Widen a @'Field' rs@ to a @'Field' ss@ given that @rs ⊆ ss@.
widenField :: (FoldRec ss ss, RMap rs, RMap ss, RecApplicative rs, RecApplicative ss, rs ⊆ ss) => Field rs -> Field ss
widenField r =
  fromMaybe (error "widenField should be provably total, isn't") $
    firstField (rreplace (fieldToRec r) (rpure Nothing))

-- |Constraint showing that @rs ⊆ ss@ and giving enough power to widen a @'CoRec' f rs@ to @'CoRec' f ss@.
type WiderCoRec rs ss = (FoldRec ss ss, RecApplicative rs, RecApplicative ss, rs ⊆ ss)

-- |Constraint on a pair of types @r@ and @s@ which are 'Wrapped' around @'CoRec' f rs@ and @'CoRec' f ss@ respectively and observes that those types can be unwrapped and the former can be widened to the latter.
type WrappedWiderCoRec f r rs s ss = (Wrapped r, Unwrapped r ~ CoRec f rs, Wrapped s, Unwrapped s ~ CoRec f ss, WiderCoRec rs ss)

-- |'WrappedWiderCoRec' specialized to 'CoRec' 'Identity', i.e. 'Field'.
type WrappedWiderField r rs s ss = WrappedWiderCoRec Identity r rs s ss

-- |Widen a @'CoRec' f rs@ wrapped in @r@ to a @'CoRec' f ss@ wrapped in @s@.
widenWrappedCoRec :: WrappedWiderCoRec f r rs s ss => r -> s
widenWrappedCoRec = review _Wrapped' . widenCoRec . view _Wrapped'

-- |Widen a @'Field' rs@ wrapped in @r@ to a @'Field' ss@ wrapped in @s@.
widenWrappedField :: WrappedWiderField r rs s ss => r -> s
widenWrappedField = review _Wrapped' . widenCoRec . view _Wrapped'

-- |Widen a wrapped 'CoRec' using 'widenWrappedCoRec' and then throw it as an 'ExceptT' error, as a convenience when using 'CoRec' / 'Field' for building up error types by set union.
throwWidened :: (Monad m, WrappedWiderCoRec f e er e' er') => e -> ExceptT e' m a
throwWidened = throwError . widenWrappedCoRec

-- |Catch a wrapped 'CoRec' thrown by an 'ExceptT' and widen it using 'widenWrappedCoRec' then rethrow it, as a convenience when using 'CoRec' / 'Field' for building up error types by set union.
rethrowWidened :: (Functor m, WrappedWiderCoRec f e er e' er') => ExceptT e m a -> ExceptT e' m a
rethrowWidened = withExceptT widenWrappedCoRec

-- |Specialized type of 'Case' for 'Val', with the type parameters in a convenient order for type application, e.g.:
--
-- @
--   valCase @"foo" (\ a -> ...)
--     :& valCase @"bar" (\ b -> ...)
--     :& RNil
-- @
valCase :: forall s a b. (a -> b) -> Case b (s :-> a)
valCase f = Case $ \(Val a) -> f a

-- |Specialized type of 'Case'' for 'Val', with the type parameters in a convenient order for type application, e.g.:
--
-- @
--   valCase' @"foo" (\ fa -> ...)
--     :& valCase' @"bar" (\ fb -> ...)
--     :& RNil
-- @
valCase' :: forall s f a b. Functor f => (f a -> b) -> Case' f b (s :-> a)
valCase' f = Case' $ \fva -> f (fmap getVal fva)

-- |Make a 'Case' which yields the symbol text for a field @s :-> ()@. E.g.:
--
-- @
--   keywordCase @"foo" (Identity (Val ())) == "foo"
-- @
keywordCase :: KnownSymbol s => Case Text (s :-> ())
keywordCase = keywordCase' id

-- |Make a 'Case'' which yields the symbol text for a field @s :-> ()@ as projected through the given function. E.g.:
--
-- @
--   keywordCase @"foo" (<> "bar") (Identity (Val ())) == "foobar"
-- @
keywordCase' :: KnownSymbol s => (Text -> a) -> Case a (s :-> ())
keywordCase' f = Case $ \(Val () :: s :-> ()) -> f (pack (symbolVal (Proxy @s)))