{-# LANGUAGE PackageImports #-}
{-# LANGUAGE AllowAmbiguousTypes   #-}
{-# LANGUAGE ConstraintKinds       #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DefaultSignatures     #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE KindSignatures        #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds             #-}
{-# LANGUAGE Rank2Types            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}

module Data.Generics.Product.Internal.Types where

import Data.Kind
import Data.Int (Int8, Int16, Int32, Int64)
import Data.Word (Word8, Word16, Word32, Word64)
import qualified Data.Text as T

import GHC.Generics
import Data.Generics.Internal.GenericN
import GHC.TypeLits
import Data.Generics.Internal.Errors

import "this" Data.Generics.Internal.VL.Traversal


-- | The children of a type are the types of its fields.
-- The 'Children' type family maps a type @a@ to its set of children.
--
-- This type family is parameterized by a symbol @ch@ (that can be declared as
-- an empty data type).
-- The symbol 'ChGeneric' provides a default definition. You can create new
-- symbols to override the set of children of abstract, non-generic types.
--
-- The following example declares a @Custom@ symbol to redefine 'Children'
-- for some abstract types from the @time@ library.
--
-- @
-- data Custom
-- type instance 'Children' Custom a = ChildrenCustom a
--
-- type family ChildrenCustom (a :: Type) where
--   ChildrenCustom DiffTime        = '[]
--   ChildrenCustom NominalDiffTime = '[]
--   -- Add more custom mappings here.
--
--   ChildrenCustom a = Children ChGeneric a
-- @
--
-- To use this definition, replace 'types' with @'typesUsing' \@Custom@.
type family Children (ch :: Type) (a :: Type) :: [Type]

-- | The default definition of 'Children'.
-- Primitive types from core libraries have no children, and other types are
-- assumed to be 'Generic'.
data ChGeneric
type instance Children ChGeneric a = ChildrenDefault a

type family ChildrenDefault (a :: Type) :: [Type] where
  ChildrenDefault Char    = '[]
  ChildrenDefault Double  = '[]
  ChildrenDefault Float   = '[]
  ChildrenDefault Integer = '[]
  ChildrenDefault Int     = '[]
  ChildrenDefault Int8    = '[]
  ChildrenDefault Int16   = '[]
  ChildrenDefault Int32   = '[]
  ChildrenDefault Int64   = '[]
  ChildrenDefault Word    = '[]
  ChildrenDefault Word8   = '[]
  ChildrenDefault Word16  = '[]
  ChildrenDefault Word32  = '[]
  ChildrenDefault Word64  = '[]
  ChildrenDefault T.Text  = '[]
  ChildrenDefault (Param n _) = '[]
  ChildrenDefault a
   = Defined (Rep a)
    (NoGeneric a
      '[ 'Text "arising from a generic traversal."
       , 'Text "Either derive the instance, or define a custom traversal using HasTypesCustom"
       ])
    (ChildrenGeneric (Rep a) '[])

type family ChildrenGeneric (f :: k -> Type) (cs :: [Type]) :: [Type] where
  ChildrenGeneric (M1 _ _ f) cs = ChildrenGeneric f cs
  ChildrenGeneric (l :*: r) cs = ChildrenGeneric l (ChildrenGeneric r cs)
  ChildrenGeneric (l :+: r) cs = ChildrenGeneric l (ChildrenGeneric r cs)
  ChildrenGeneric (Rec0 a) cs = a ': cs
  ChildrenGeneric _ cs = cs

type Interesting (ch :: Type) (a :: Type) (t :: Type)
  = Defined_list (Children ch t) (NoChildren ch t)
    (IsNothing (Interesting' ch a '[t] (Children ch t)))

type family NoChildren (ch :: Type) (a :: Type) :: Constraint where
  NoChildren ch a = PrettyError
                    '[ 'Text "No type family instance for " ':<>: QuoteType (Children ch a)
                     , 'Text "arising from a traversal over " ':<>: QuoteType a
                     , 'Text "with custom strategy " ':<>: QuoteType ch
                     ]
                    

type family Interesting' (ch :: Type) (a :: Type) (seen :: [Type]) (ts :: [Type]) :: Maybe [Type] where
  Interesting' ch _ seen '[] = 'Just seen
  Interesting' ch a seen (t ': ts) =
    InterestingOr ch a (InterestingUnless ch a seen t (Elem t seen)) ts

-- Short circuit
-- Note: we only insert 't' to the seen list if it's not already there (which is precisely when `s` is 'False)
type family InterestingUnless
    (ch :: Type) (a :: Type) (seen :: [Type]) (t :: Type) (alreadySeen :: Bool) ::
    Maybe [Type] where
  InterestingUnless ch a seen a _ = 'Nothing
  InterestingUnless ch a seen t 'True = 'Just seen
  InterestingUnless ch a seen t 'False
    = Defined_list (Children ch t) (NoChildren ch t)
      (Interesting' ch a (t ': seen) (Children ch t))

-- Short circuit
type family InterestingOr
    (ch :: Type) (a :: Type) (seen' :: Maybe [Type]) (ts :: [Type]) ::
    Maybe [Type] where
  InterestingOr ch a 'Nothing _ = 'Nothing
  InterestingOr ch a ('Just seen) ts = Interesting' ch a seen ts

type family Elem a as where
  Elem a (a ': _) = 'True
  Elem a (_ ': as) = Elem a as
  Elem a '[] = 'False

type family IsNothing a where
  IsNothing ('Just _) = 'False
  IsNothing 'Nothing = 'True

--------------------------------------------------------------------------------

--------------------------------------------------------------------------------
-- HasTypes
--------------------------------------------------------------------------------

class HasTypes s a where
  types_ :: Traversal' s a

  types_ a -> f a
_ = s -> f s
forall (f :: * -> *) a. Applicative f => a -> f a
pure
  {-# INLINE types_ #-}

instance
  ( HasTypesUsing ChGeneric s s a a
  ) => HasTypes s a where
  types_ :: (a -> f a) -> s -> f s
types_ = forall s t a b.
HasTypesUsing ChGeneric s t a b =>
Traversal s t a b
forall ch s t a b. HasTypesUsing ch s t a b => Traversal s t a b
typesUsing_ @ChGeneric
  {-# INLINE types_ #-}

--------------------------------------------------------------------------------
data Void
instance {-# OVERLAPPING #-} HasTypes Void a where
  types_ :: (a -> f a) -> Void -> f Void
types_ a -> f a
_ = Void -> f Void
forall (f :: * -> *) a. Applicative f => a -> f a
pure

instance {-# OVERLAPPING #-} HasTypes s Void where
  types_ :: (Void -> f Void) -> s -> f s
types_ Void -> f Void
_ = s -> f s
forall (f :: * -> *) a. Applicative f => a -> f a
pure

instance {-# OVERLAPPING #-} HasTypesUsing ch Void Void a b where
  typesUsing_ :: (a -> f b) -> Void -> f Void
typesUsing_ a -> f b
_ = Void -> f Void
forall (f :: * -> *) a. Applicative f => a -> f a
pure

instance {-# OVERLAPPING #-} HasTypesUsing ch s s Void Void where
  typesUsing_ :: (Void -> f Void) -> s -> f s
typesUsing_ Void -> f Void
_ = s -> f s
forall (f :: * -> *) a. Applicative f => a -> f a
pure
--------------------------------------------------------------------------------

--------------------------------------------------------------------------------
-- HasTypesUsing
--------------------------------------------------------------------------------

-- | @since 1.2.0.0
class HasTypesUsing (ch :: Type) s t a b where
  typesUsing_ :: Traversal s t a b

instance {-# OVERLAPPABLE #-}
  ( HasTypesOpt ch (Interesting ch a s) s t a b
  ) => HasTypesUsing ch s t a b where
  typesUsing_ :: (a -> f b) -> s -> f t
typesUsing_ = forall ch (p :: Bool) s t a b.
HasTypesOpt ch p s t a b =>
Traversal s t a b
forall s t a b.
HasTypesOpt ch (Interesting ch a s) s t a b =>
Traversal s t a b
typesOpt @ch @(Interesting ch a s)
  {-# INLINE typesUsing_ #-}

instance {-# OVERLAPPABLE #-} HasTypesUsing ch a b a b where
  typesUsing_ :: (a -> f b) -> a -> f b
typesUsing_ = (a -> f b) -> a -> f b
forall a. a -> a
id

-- | By adding instances to this class, we can override the default
-- behaviour in an ad-hoc manner.
-- For example:
--
-- @
-- instance HasTypesCustom Custom Opaque Opaque String String where
--   typesCustom f (Opaque str) = Opaque <$> f str
-- @
--
-- @since 1.2.0.0
class HasTypesCustom (ch :: Type) s t a b where
  -- | This function should never be used directly, only to override
  -- the default traversal behaviour. To actually use the custom
  -- traversal strategy, see 'typesUsing'. This is because 'typesUsing' does
  -- additional optimisations, like ensuring that nodes with no relevant members will
  -- not be traversed at runtime.
  typesCustom :: Traversal s t a b

instance {-# OVERLAPPABLE #-}
  ( GHasTypes ch (Rep s) (Rep t) a b
  , Generic s
  , Generic t
  -- if there's no Generic instance here, it means we got through the
  -- Children check by a user-defined custom strategy.
  -- Therefore, we can ignore the missing Generic instance, and
  -- instead report a missing HasTypesCustom instance
  , Defined (Rep s)
    (PrettyError '[ 'Text "No instance " ':<>: QuoteType (HasTypesCustom ch s t a b)])
    (() :: Constraint)
  ) => HasTypesCustom ch s t a b where
  typesCustom :: (a -> f b) -> s -> f t
typesCustom a -> f b
f s
s = Rep t Any -> t
forall a x. Generic a => Rep a x -> a
to (Rep t Any -> t) -> f (Rep t Any) -> f t
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> f b) -> Rep s Any -> f (Rep t Any)
forall k k (ch :: k) (s :: k -> *) (t :: k -> *) a b (x :: k).
GHasTypes ch s t a b =>
Traversal (s x) (t x) a b
gtypes_ @ch a -> f b
f (s -> Rep s Any
forall a x. Generic a => a -> Rep a x
from s
s)

--------------------------------------------------------------------------------
-- Internals
--------------------------------------------------------------------------------
-- TODO: these should never leak out in error messages

class HasTypesOpt (ch :: Type) (p :: Bool) s t a b where
  typesOpt :: Traversal s t a b

instance HasTypesCustom ch s t a b => HasTypesOpt ch 'True s t a b where
  typesOpt :: (a -> f b) -> s -> f t
typesOpt = forall s t a b. HasTypesCustom ch s t a b => Traversal s t a b
forall ch s t a b. HasTypesCustom ch s t a b => Traversal s t a b
typesCustom @ch

instance HasTypesOpt ch 'False s s a b where
  typesOpt :: (a -> f b) -> s -> f s
typesOpt a -> f b
_ = s -> f s
forall (f :: * -> *) a. Applicative f => a -> f a
pure

--------------------------------------------------------------------------------
-- TODO: pull out recursion here.

class GHasTypes ch s t a b where
  gtypes_ :: Traversal (s x) (t x) a b

instance
  ( GHasTypes ch l l' a b
  , GHasTypes ch r r' a b
  ) => GHasTypes ch (l :*: r) (l' :*: r') a b where
  gtypes_ :: (a -> f b) -> (:*:) l r x -> f ((:*:) l' r' x)
gtypes_ a -> f b
f (l x
l :*: r x
r) = l' x -> r' x -> (:*:) l' r' x
forall k (f :: k -> *) (g :: k -> *) (p :: k).
f p -> g p -> (:*:) f g p
(:*:) (l' x -> r' x -> (:*:) l' r' x)
-> f (l' x) -> f (r' x -> (:*:) l' r' x)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> f b) -> l x -> f (l' x)
forall k k (ch :: k) (s :: k -> *) (t :: k -> *) a b (x :: k).
GHasTypes ch s t a b =>
Traversal (s x) (t x) a b
gtypes_ @ch a -> f b
f l x
l f (r' x -> (:*:) l' r' x) -> f (r' x) -> f ((:*:) l' r' x)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (a -> f b) -> r x -> f (r' x)
forall k k (ch :: k) (s :: k -> *) (t :: k -> *) a b (x :: k).
GHasTypes ch s t a b =>
Traversal (s x) (t x) a b
gtypes_ @ch a -> f b
f r x
r
  {-# INLINE gtypes_ #-}

instance
  ( GHasTypes ch l l' a b
  , GHasTypes ch r r' a b
  ) => GHasTypes ch (l :+: r) (l' :+: r') a b where
  gtypes_ :: (a -> f b) -> (:+:) l r x -> f ((:+:) l' r' x)
gtypes_ a -> f b
f (L1 l x
l) = l' x -> (:+:) l' r' x
forall k (f :: k -> *) (g :: k -> *) (p :: k). f p -> (:+:) f g p
L1 (l' x -> (:+:) l' r' x) -> f (l' x) -> f ((:+:) l' r' x)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> f b) -> l x -> f (l' x)
forall k k (ch :: k) (s :: k -> *) (t :: k -> *) a b (x :: k).
GHasTypes ch s t a b =>
Traversal (s x) (t x) a b
gtypes_ @ch a -> f b
f l x
l
  gtypes_ a -> f b
f (R1 r x
r) = r' x -> (:+:) l' r' x
forall k (f :: k -> *) (g :: k -> *) (p :: k). g p -> (:+:) f g p
R1 (r' x -> (:+:) l' r' x) -> f (r' x) -> f ((:+:) l' r' x)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> f b) -> r x -> f (r' x)
forall k k (ch :: k) (s :: k -> *) (t :: k -> *) a b (x :: k).
GHasTypes ch s t a b =>
Traversal (s x) (t x) a b
gtypes_ @ch a -> f b
f r x
r
  {-# INLINE gtypes_ #-}

instance GHasTypes ch s t a b => GHasTypes ch (M1 m meta s) (M1 m meta t) a b where
  gtypes_ :: (a -> f b) -> M1 m meta s x -> f (M1 m meta t x)
gtypes_ a -> f b
f (M1 s x
s) = t x -> M1 m meta t x
forall k i (c :: Meta) (f :: k -> *) (p :: k). f p -> M1 i c f p
M1 (t x -> M1 m meta t x) -> f (t x) -> f (M1 m meta t x)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> f b) -> s x -> f (t x)
forall k k (ch :: k) (s :: k -> *) (t :: k -> *) a b (x :: k).
GHasTypes ch s t a b =>
Traversal (s x) (t x) a b
gtypes_ @ch a -> f b
f s x
s
  {-# INLINE gtypes_ #-}

-- In the recursive case, we invoke 'HasTypesUsing' again, using the
-- same strategy
-- This instance is marked INCOHERENT, because
instance {-# INCOHERENT #-} HasTypesUsing ch s t a b => GHasTypes ch (Rec0 s) (Rec0 t) a b where
  gtypes_ :: (a -> f b) -> Rec0 s x -> f (Rec0 t x)
gtypes_ a -> f b
f (K1 s
x) = t -> Rec0 t x
forall k i c (p :: k). c -> K1 i c p
K1 (t -> Rec0 t x) -> f t -> f (Rec0 t x)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> f b) -> s -> f t
forall ch s t a b. HasTypesUsing ch s t a b => Traversal s t a b
typesUsing_ @ch a -> f b
f s
x
  {-# INLINE gtypes_ #-}

-- | The default instance for 'HasTypes' acts as a synonym for
-- 'HasTypesUsing ChGeneric', so in most cases this instance should
-- behave the same as the one above.
-- However, there might be overlapping instances defined for
-- 'HasTypes' directly, in which case we want to prefer those
-- instances (even though the custom instances should always be added to 'HasTypesCustom')
instance {-# OVERLAPPING #-} HasTypes b a => GHasTypes ChGeneric (Rec0 b) (Rec0 b) a a where
  gtypes_ :: (a -> f a) -> Rec0 b x -> f (Rec0 b x)
gtypes_ a -> f a
f (K1 b
x) = b -> Rec0 b x
forall k i c (p :: k). c -> K1 i c p
K1 (b -> Rec0 b x) -> f b -> f (Rec0 b x)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (a -> f a) -> b -> f b
forall s a. HasTypes s a => Traversal' s a
types_ @b @a a -> f a
f b
x
  {-# INLINE gtypes_ #-}

instance GHasTypes ch U1 U1 a b where
  gtypes_ :: (a -> f b) -> U1 x -> f (U1 x)
gtypes_ a -> f b
_ U1 x
_ = U1 x -> f (U1 x)
forall (f :: * -> *) a. Applicative f => a -> f a
pure U1 x
forall k (p :: k). U1 p
U1
  {-# INLINE gtypes_ #-}

instance GHasTypes ch V1 V1 a b where
  gtypes_ :: (a -> f b) -> V1 x -> f (V1 x)
gtypes_ a -> f b
_ = V1 x -> f (V1 x)
forall (f :: * -> *) a. Applicative f => a -> f a
pure
  {-# INLINE gtypes_ #-}