{-# options_ghc -Wall #-}
{-# language TemplateHaskellQuotes #-}
{-# language FlexibleContexts #-}
{-# language FlexibleInstances #-}
{-# language ScopedTypeVariables #-}
{-# language KindSignatures #-}
{-# language DataKinds #-}
{-# language TypeApplications #-}
{-# language TypeOperators #-}
{-# language EmptyCase #-}
{-# language DefaultSignatures #-}

module Language.Haskell.TH.TraverseCode
  ( TraverseCode (..)
  , sequenceCode
  , genericTraverseCode
  , genericSequenceCode
  ) where
import Generics.Linear
import Language.Haskell.TH.Syntax (Code, Lift (..), Exp (..), Quote, Name)
import qualified Language.Haskell.TH.Syntax as TH
import Language.Haskell.TH.Lib (conE)
import Data.Kind (Type)

-- for instances
import qualified Data.Functor.Product as FProd
import qualified Data.Functor.Sum as FSum
import Data.Functor.Identity
import qualified Data.Sequence.Internal as Seq
import Data.Coerce
import Control.Applicative

-- | Containers supporting \"traversal\" in 'Code'.
class TraverseCode t where
  -- | Given a container and a function to fill it with splices,
  -- produce a splice that will generate a container of their results.
  traverseCode :: Quote m => (a -> Code m b) -> t a -> Code m (t b)

  default traverseCode :: (Quote m, GTraverseCode (Rep1 t), Generic1 t) => (a -> Code m b) -> t a -> Code m (t b)
  traverseCode = (a -> Code m b) -> t a -> Code m (t b)
forall (m :: * -> *) (t :: * -> *) a b.
(Quote m, GTraverseCode (Rep1 t), Generic1 t) =>
(a -> Code m b) -> t a -> Code m (t b)
genericTraverseCode

-- | Given a container of splices, produce a splice that will generate a
-- container of their results.
sequenceCode :: (TraverseCode t, Quote m) => t (Code m a) -> Code m (t a)
sequenceCode :: forall (t :: * -> *) (m :: * -> *) a.
(TraverseCode t, Quote m) =>
t (Code m a) -> Code m (t a)
sequenceCode = (Code m a -> Code m a) -> t (Code m a) -> Code m (t a)
forall (t :: * -> *) (m :: * -> *) a b.
(TraverseCode t, Quote m) =>
(a -> Code m b) -> t a -> Code m (t b)
traverseCode Code m a -> Code m a
forall a. a -> a
id

-- | Like 'sequenceCode', but using the @"Generics.Linear".'Generic1'@ instance.
genericSequenceCode :: (Quote m, GTraverseCode (Rep1 t), Generic1 t) => t (Code m a) -> Code m (t a)
genericSequenceCode :: forall (m :: * -> *) (t :: * -> *) a.
(Quote m, GTraverseCode (Rep1 t), Generic1 t) =>
t (Code m a) -> Code m (t a)
genericSequenceCode = m Exp -> Code m (t a)
forall a (m :: * -> *). Quote m => m Exp -> Code m a
TH.unsafeCodeCoerce (m Exp -> Code m (t a))
-> (t (Code m a) -> m Exp) -> t (Code m a) -> Code m (t a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Code m a -> Code m a) -> Rep1 t (Code m a) -> m Exp
forall (f :: * -> *) (m :: * -> *) a b.
(GTraverseCode f, Quote m) =>
(a -> Code m b) -> f a -> m Exp
gtraverseCode Code m a -> Code m a
forall a. a -> a
id (Rep1 t (Code m a) -> m Exp)
-> (t (Code m a) -> Rep1 t (Code m a)) -> t (Code m a) -> m Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. t (Code m a) -> Rep1 t (Code m a)
forall {k} (f :: k -> *) (p :: k). Generic1 f => f p -> Rep1 f p
from1

-- | Like 'traverseCode', but using the @"Generics.Linear".'Generic1'@ instance.
genericTraverseCode :: (Quote m, GTraverseCode (Rep1 t), Generic1 t) => (a -> Code m b) -> t a -> Code m (t b)
genericTraverseCode :: forall (m :: * -> *) (t :: * -> *) a b.
(Quote m, GTraverseCode (Rep1 t), Generic1 t) =>
(a -> Code m b) -> t a -> Code m (t b)
genericTraverseCode a -> Code m b
f = m Exp -> Code m (t b)
forall a (m :: * -> *). Quote m => m Exp -> Code m a
TH.unsafeCodeCoerce (m Exp -> Code m (t b)) -> (t a -> m Exp) -> t a -> Code m (t b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Code m b) -> Rep1 t a -> m Exp
forall (f :: * -> *) (m :: * -> *) a b.
(GTraverseCode f, Quote m) =>
(a -> Code m b) -> f a -> m Exp
gtraverseCode a -> Code m b
f (Rep1 t a -> m Exp) -> (t a -> Rep1 t a) -> t a -> m Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. t a -> Rep1 t a
forall {k} (f :: k -> *) (p :: k). Generic1 f => f p -> Rep1 f p
from1

class GTraverseCode f where
  gtraverseCode :: Quote m => (a -> Code m b) -> f a -> m Exp

data Goop (d :: Meta) (f :: Type -> Type) a = Goop

instance (Datatype c, GTraverseCodeCon f) => GTraverseCode (D1 c f) where
  gtraverseCode :: forall (m :: * -> *) a b.
Quote m =>
(a -> Code m b) -> D1 c f a -> m Exp
gtraverseCode a -> Code m b
f (M1 f a
x) = String -> String -> (a -> Code m b) -> f a -> m Exp
forall (f :: * -> *) (m :: * -> *) a b.
(GTraverseCodeCon f, Quote m) =>
String -> String -> (a -> Code m b) -> f a -> m Exp
gtraverseCodeCon String
pkg String
modl a -> Code m b
f f a
x
    where
      pkg :: String
pkg = Goop c f Any -> String
forall {k} (d :: k) k1 (t :: k -> (k1 -> *) -> k1 -> *)
       (f :: k1 -> *) (a :: k1).
Datatype d =>
t d f a -> String
packageName (forall (d :: Meta) (f :: * -> *) a. Goop d f a
Goop @c @f)
      modl :: String
modl = Goop c f Any -> String
forall {k} (d :: k) k1 (t :: k -> (k1 -> *) -> k1 -> *)
       (f :: k1 -> *) (a :: k1).
Datatype d =>
t d f a -> String
moduleName (forall (d :: Meta) (f :: * -> *) a. Goop d f a
Goop @c @f)

class GTraverseCodeCon f where
  gtraverseCodeCon :: Quote m => String -> String -> (a -> Code m b) -> f a -> m Exp

-- This instance seems totally useless, but it's obviously valid.
instance GTraverseCodeCon V1 where
  gtraverseCodeCon :: forall (m :: * -> *) a b.
Quote m =>
String -> String -> (a -> Code m b) -> V1 a -> m Exp
gtraverseCodeCon String
_pkg String
_modl a -> Code m b
_f V1 a
x = case V1 a
x of

instance (GTraverseCodeCon f, GTraverseCodeCon g) => GTraverseCodeCon (f :+: g) where
  gtraverseCodeCon :: forall (m :: * -> *) a b.
Quote m =>
String -> String -> (a -> Code m b) -> (:+:) f g a -> m Exp
gtraverseCodeCon String
pkg String
modl a -> Code m b
f (L1 f a
x) = String -> String -> (a -> Code m b) -> f a -> m Exp
forall (f :: * -> *) (m :: * -> *) a b.
(GTraverseCodeCon f, Quote m) =>
String -> String -> (a -> Code m b) -> f a -> m Exp
gtraverseCodeCon String
pkg String
modl a -> Code m b
f f a
x
  gtraverseCodeCon String
pkg String
modl a -> Code m b
f (R1 g a
y) = String -> String -> (a -> Code m b) -> g a -> m Exp
forall (f :: * -> *) (m :: * -> *) a b.
(GTraverseCodeCon f, Quote m) =>
String -> String -> (a -> Code m b) -> f a -> m Exp
gtraverseCodeCon String
pkg String
modl a -> Code m b
f g a
y

instance (Constructor c, GTraverseCodeFields f) => GTraverseCodeCon (C1 c f) where
  gtraverseCodeCon :: forall (m :: * -> *) a b.
Quote m =>
String -> String -> (a -> Code m b) -> C1 c f a -> m Exp
gtraverseCodeCon String
pkg String
modl a -> Code m b
f (M1 f a
x) = m Exp -> (a -> Code m b) -> f a -> m Exp
forall (f :: * -> *) (m :: * -> *) a b.
(GTraverseCodeFields f, Quote m) =>
m Exp -> (a -> Code m b) -> f a -> m Exp
gtraverseCodeFields (Name -> m Exp
forall (m :: * -> *). Quote m => Name -> m Exp
conE Name
conN) a -> Code m b
f f a
x
    where
      conBase :: String
conBase = Goop c f Any -> String
forall {k} (c :: k) k1 (t :: k -> (k1 -> *) -> k1 -> *)
       (f :: k1 -> *) (a :: k1).
Constructor c =>
t c f a -> String
conName (forall (d :: Meta) (f :: * -> *) a. Goop d f a
Goop @c @f)
      conN :: Name
      conN :: Name
conN = String -> String -> String -> Name
TH.mkNameG_d String
pkg String
modl String
conBase

class GTraverseCodeFields f where
  gtraverseCodeFields :: Quote m => m Exp -> (a -> Code m b) -> f a -> m Exp

instance GTraverseCodeFields f => GTraverseCodeFields (S1 c f) where
  gtraverseCodeFields :: forall (m :: * -> *) a b.
Quote m =>
m Exp -> (a -> Code m b) -> S1 c f a -> m Exp
gtraverseCodeFields m Exp
c a -> Code m b
f (M1 f a
x) = m Exp -> (a -> Code m b) -> f a -> m Exp
forall (f :: * -> *) (m :: * -> *) a b.
(GTraverseCodeFields f, Quote m) =>
m Exp -> (a -> Code m b) -> f a -> m Exp
gtraverseCodeFields m Exp
c a -> Code m b
f f a
x

instance (GTraverseCodeFields f, GTraverseCodeFields g) => GTraverseCodeFields (f :*: g) where
  gtraverseCodeFields :: forall (m :: * -> *) a b.
Quote m =>
m Exp -> (a -> Code m b) -> (:*:) f g a -> m Exp
gtraverseCodeFields m Exp
c a -> Code m b
f (f a
x :*: g a
y) =
    m Exp -> (a -> Code m b) -> g a -> m Exp
forall (f :: * -> *) (m :: * -> *) a b.
(GTraverseCodeFields f, Quote m) =>
m Exp -> (a -> Code m b) -> f a -> m Exp
gtraverseCodeFields (m Exp -> (a -> Code m b) -> f a -> m Exp
forall (f :: * -> *) (m :: * -> *) a b.
(GTraverseCodeFields f, Quote m) =>
m Exp -> (a -> Code m b) -> f a -> m Exp
gtraverseCodeFields m Exp
c a -> Code m b
f f a
x) a -> Code m b
f g a
y

instance Lift p => GTraverseCodeFields (K1 i p) where
  gtraverseCodeFields :: forall (m :: * -> *) a b.
Quote m =>
m Exp -> (a -> Code m b) -> K1 i p a -> m Exp
gtraverseCodeFields m Exp
c a -> Code m b
_f (K1 p
x) = [| $c x |]

instance GTraverseCodeFields Par1 where
  gtraverseCodeFields :: forall (m :: * -> *) a b.
Quote m =>
m Exp -> (a -> Code m b) -> Par1 a -> m Exp
gtraverseCodeFields m Exp
cc a -> Code m b
f (Par1 a
ca) = [| $cc $(TH.unTypeCode (f ca)) |]

instance GTraverseCodeFields U1 where
  gtraverseCodeFields :: forall (m :: * -> *) a b.
Quote m =>
m Exp -> (a -> Code m b) -> U1 a -> m Exp
gtraverseCodeFields m Exp
cc a -> Code m b
_f U1 a
U1 = m Exp
cc

instance (GTraverseCodeFields f, TraverseCode g) => GTraverseCodeFields (f :.: g) where
  gtraverseCodeFields :: forall (m :: * -> *) a b.
Quote m =>
m Exp -> (a -> Code m b) -> (:.:) f g a -> m Exp
gtraverseCodeFields m Exp
cc a -> Code m b
f (Comp1 f (g a)
x) =
    m Exp -> (g a -> Code m (g b)) -> f (g a) -> m Exp
forall (f :: * -> *) (m :: * -> *) a b.
(GTraverseCodeFields f, Quote m) =>
m Exp -> (a -> Code m b) -> f a -> m Exp
gtraverseCodeFields m Exp
cc ((a -> Code m b) -> g a -> Code m (g b)
forall (t :: * -> *) (m :: * -> *) a b.
(TraverseCode t, Quote m) =>
(a -> Code m b) -> t a -> Code m (t b)
traverseCode a -> Code m b
f) f (g a)
x

-- TraverseCode instances

instance TraverseCode Maybe
instance TraverseCode Identity
instance TraverseCode []
instance TH.Lift a => TraverseCode (Either a)
instance TH.Lift a => TraverseCode ((,) a)
instance (TraverseCode f, TraverseCode g) => TraverseCode (FProd.Product f g)
instance (TraverseCode f, TraverseCode g) => TraverseCode (FSum.Sum f g)
instance Lift a => TraverseCode (Const a)

instance TraverseCode V1
instance TraverseCode U1
instance (TraverseCode f, TraverseCode g) => TraverseCode (f :*: g)
instance (TraverseCode f, TraverseCode g) => TraverseCode (f :+: g)
instance TraverseCode f => TraverseCode (M1 i c f)
instance TraverseCode Par1
instance Lift a => TraverseCode (K1 i a)

-- The Elem instance isn't needed for the Seq instance
instance TraverseCode Seq.Elem
instance TraverseCode Seq.Digit
instance TraverseCode Seq.Node
instance TraverseCode Seq.FingerTree
instance TraverseCode Seq.Seq where
  -- This wonky way of doing it makes for a more compact
  -- splice.
  traverseCode :: forall (m :: * -> *) a b.
Quote m =>
(a -> Code m b) -> Seq a -> Code m (Seq b)
traverseCode a -> Code m b
f Seq a
s = [|| coerceFT $$(traverseCode f ft') ||]
    where
      ft' :: FingerTree a
ft' = Seq a -> FingerTree a
forall a. Seq a -> FingerTree a
coerceSeq ((() -> a -> a) -> Seq () -> Seq a -> Seq a
forall a b c. (a -> b -> c) -> Seq a -> Seq b -> Seq c
Seq.zipWith ((a -> () -> a) -> () -> a -> a
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> () -> a
forall a b. a -> b -> a
const) (Int -> () -> Seq ()
forall a. Int -> a -> Seq a
Seq.replicate (Seq a -> Int
forall a. Seq a -> Int
Seq.length Seq a
s) ()) Seq a
s)
coerceFT :: Seq.FingerTree a -> Seq.Seq a
coerceFT :: forall a. FingerTree a -> Seq a
coerceFT = FingerTree a -> Seq a
coerce
coerceSeq :: Seq.Seq a -> Seq.FingerTree a
coerceSeq :: forall a. Seq a -> FingerTree a
coerceSeq = Seq a -> FingerTree a
coerce