{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
module Data.Vinyl.XRec where
import Data.Vinyl.Core (Rec(..))
import Data.Vinyl.Functor
import Data.Vinyl.Lens (RecElem, RecElemFCtx, rgetC)
import Data.Vinyl.TypeLevel (RIndex)
import Data.Monoid
import GHC.TypeLits (KnownSymbol)
type XRec f = Rec (XData f)
pattern (::&) :: HKD f r -> XRec f rs -> XRec f (r ': rs)
pattern x ::& xs = XData x :& xs
{-# COMPLETE (::&) #-}
infixr 7 ::&
pattern XRNil :: XRec f '[]
pattern XRNil = RNil
{-# COMPLETE XRNil #-}
rmapX :: forall f g rs. (XRMap f g rs, IsoXRec f rs, IsoXRec g rs)
      => (forall a. HKD f a -> HKD g a) -> Rec f rs -> Rec g rs
rmapX f = fromXRec . xrmapAux aux . toXRec
  where aux :: forall a. XData f a -> XData g a
        aux = XData . f @a . unX
rmapXEndo :: forall f rs. (XRMap f f rs, IsoXRec f rs)
          => (forall a. HKD f a -> HKD f a) -> Rec f rs -> Rec f rs
rmapXEndo f = fromXRec . xrmapAux aux . toXRec
  where aux :: forall a. XData f a -> XData f a
        aux = XData . f @a . unX
xrmap :: forall f g rs. XRMap f g rs
      => (forall a. HKD f a -> HKD g a) -> XRec f rs -> XRec g rs
xrmap f = xrmapAux aux
  where aux :: forall a. XData f a -> XData g a
        aux = XData . f @a . unX
newtype XData t a = XData { unX :: HKD t a }
class XRMap f g rs where
  xrmapAux :: (forall a . XData f a -> XData g a) -> XRec f rs -> XRec g rs
instance XRMap f g '[] where
  xrmapAux _ RNil = RNil
instance forall f g r rs. (XRMap f g rs, IsoHKD f r, IsoHKD g r)
  => XRMap f g (r ': rs) where
  xrmapAux f (x :& xs) = f x :& xrmapAux f xs
class XRApply f g rs where
  xrapply :: XRec (Lift (->) f g) rs -> XRec f rs -> XRec g rs
instance XRApply f g '[] where
  xrapply RNil RNil = RNil
instance XRApply f g rs => XRApply f g (r ': rs) where
  xrapply (XData f :& fs) (XData x :& xs) = XData (f x) :& xrapply fs xs
class IsoXRec f ts where
  fromXRec :: XRec f ts -> Rec f ts
  toXRec :: Rec f ts -> XRec f ts
instance IsoXRec f '[] where
  fromXRec RNil = RNil
  toXRec RNil = XRNil
instance (IsoXRec f ts, IsoHKD f t) => IsoXRec f (t ': ts) where
  fromXRec (x ::& xs) = unHKD x :& fromXRec xs
  toXRec (x :& xs) = toHKD x ::& toXRec xs
class IsoHKD f a where
  type HKD f a
  type HKD f a = f a
  unHKD :: HKD f a -> f a
  default unHKD :: HKD f a ~ f a => HKD f a -> f a
  unHKD = id
  toHKD :: f a -> HKD f a
  default toHKD :: (HKD f a ~ f a) => f a -> HKD f a
  toHKD = id
instance IsoHKD Identity a where
  type HKD Identity a = a
  unHKD = Identity
  toHKD (Identity x) = x
instance KnownSymbol s => IsoHKD ElField '(s,a) where
  type HKD ElField '(s,a) = a
  unHKD = Field
  toHKD (Field x) = x
instance (IsoHKD f (HKD g a), IsoHKD g a, Functor f) => IsoHKD (Compose f g) a where
  type HKD (Compose f g) a = HKD f (HKD g a)
  unHKD x = Compose (unHKD <$> unHKD x)
  toHKD (Compose fgx) = toHKD (toHKD <$> fgx)
instance (IsoHKD f a, IsoHKD g a) => IsoHKD (Lift (->) f g) a where
  type HKD (Lift (->) f g) a = HKD f a -> HKD g a
  unHKD x = Lift (unHKD . x . toHKD)
  toHKD (Lift x) = toHKD . x . unHKD
instance IsoHKD IO a where
instance IsoHKD (Either a) b where
instance IsoHKD Maybe a where
instance IsoHKD First a where
instance IsoHKD Last a where
instance IsoHKD ((,) a) b where
instance IsoHKD Sum a where
  type HKD Sum a = a
  unHKD = Sum
  toHKD (Sum x) = x
instance IsoHKD Product a where
  type HKD Product a = a
  unHKD = Product
  toHKD (Product x) = x
rgetX :: forall a record f rs.
         (RecElem record a a rs rs (RIndex a rs),
          RecElemFCtx record f,
          IsoHKD f a)
      => record f rs -> HKD f a
rgetX = toHKD . rgetAux @a
  where rgetAux :: forall r.
                   (RecElem record r r rs rs (RIndex r rs),
                    RecElemFCtx record f)
                => record f rs -> f r
        rgetAux = rgetC