{-# language TemplateHaskellQuotes #-}
{-# language FlexibleContexts #-}
{-# language FlexibleInstances #-}
{-# language ScopedTypeVariables #-}
{-# language DataKinds #-}
{-# language TypeOperators #-}
{-# language EmptyCase #-}
{-# language DefaultSignatures #-}
{-# language BangPatterns #-}
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 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
import qualified Data.Semigroup as S
import qualified Data.Monoid as M
import qualified Data.Complex as Complex
import qualified Data.Ord as Ord
import qualified Data.Functor.Compose as Compose
import qualified Data.Tree as Tree
import qualified Data.List.NonEmpty as NonEmpty
import GHC.Tuple (Solo)
import Data.Proxy (Proxy)
import qualified GHC.Generics as GHCGenerics
import qualified Data.Array as Ar
import qualified Data.Primitive.Array as PAr
import qualified Data.Primitive.SmallArray as PSmAr
import Data.Foldable (toList)
class TraverseCode t where
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
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
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
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
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 m :: D1 c f a
m@(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 = D1 c f a -> String
forall {k} (d :: k) k1 (t :: k -> (k1 -> *) -> k1 -> *)
(f :: k1 -> *) (a :: k1).
Datatype d =>
t d f a -> String
packageName D1 c f a
m
modl :: String
modl = D1 c f a -> String
forall {k} (d :: k) k1 (t :: k -> (k1 -> *) -> k1 -> *)
(f :: k1 -> *) (a :: k1).
Datatype d =>
t d f a -> String
moduleName D1 c f a
m
class GTraverseCodeCon f where
gtraverseCodeCon :: Quote m => String -> String -> (a -> Code m b) -> f a -> m Exp
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 m :: C1 c f a
m@(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 = C1 c f a -> String
forall {k} (c :: k) k1 (t :: k -> (k1 -> *) -> k1 -> *)
(f :: k1 -> *) (a :: k1).
Constructor c =>
t c f a -> String
conName C1 c f a
m
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
instance GTraverseCodeFields (URec Char) where
gtraverseCodeFields :: forall (m :: * -> *) a b.
Quote m =>
m Exp -> (a -> Code m b) -> URec Char a -> m Exp
gtraverseCodeFields m Exp
c a -> Code m b
_f (UChar Char#
ch) = [| $c (UChar ch) |]
instance GTraverseCodeFields (URec Float) where
gtraverseCodeFields :: forall (m :: * -> *) a b.
Quote m =>
m Exp -> (a -> Code m b) -> URec Float a -> m Exp
gtraverseCodeFields m Exp
c a -> Code m b
_f (UFloat Float#
ch) = [| $c (UFloat ch) |]
instance GTraverseCodeFields (URec Double) where
gtraverseCodeFields :: forall (m :: * -> *) a b.
Quote m =>
m Exp -> (a -> Code m b) -> URec Double a -> m Exp
gtraverseCodeFields m Exp
c a -> Code m b
_f (UDouble Double#
ch) = [| $c (UDouble ch) |]
instance GTraverseCodeFields (URec Int) where
gtraverseCodeFields :: forall (m :: * -> *) a b.
Quote m =>
m Exp -> (a -> Code m b) -> URec Int a -> m Exp
gtraverseCodeFields m Exp
c a -> Code m b
_f (UInt Int#
ch) = [| $c (UInt ch) |]
instance GTraverseCodeFields (URec Word) where
gtraverseCodeFields :: forall (m :: * -> *) a b.
Quote m =>
m Exp -> (a -> Code m b) -> URec Word a -> m Exp
gtraverseCodeFields m Exp
c a -> Code m b
_f (UWord Word#
ch) = [| $c (UWord ch) |]
instance GTraverseCodeFields f => GTraverseCodeFields (MP1 m f) where
gtraverseCodeFields :: forall (m :: * -> *) a b.
Quote m =>
m Exp -> (a -> Code m b) -> MP1 m f a -> m Exp
gtraverseCodeFields m Exp
c a -> Code m b
f (MP1 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 TraverseCode Maybe
instance TraverseCode Identity
instance TraverseCode []
instance TH.Lift a => TraverseCode (Either a)
instance TH.Lift a => TraverseCode ((,) a)
instance (TH.Lift a, TH.Lift b) => TraverseCode ((,,) a b)
instance (TH.Lift a, TH.Lift b, TH.Lift c) => TraverseCode ((,,,) a b c)
instance (TH.Lift a, TH.Lift b, TH.Lift c, TH.Lift d) => TraverseCode ((,,,,) a b c d)
instance (TH.Lift a, TH.Lift b, TH.Lift c, TH.Lift d, TH.Lift e) => TraverseCode ((,,,,,) a b c d e)
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)
instance (TraverseCode f, TraverseCode g) => TraverseCode (f :.: g)
instance (TraverseCode f, TraverseCode g) => TraverseCode (Compose.Compose f g)
instance TraverseCode f => TraverseCode (MP1 m f)
instance TraverseCode f => TraverseCode (GHCGenerics.Rec1 f)
instance (TraverseCode f, TraverseCode g) => TraverseCode (f GHCGenerics.:.: g)
instance TraverseCode Seq.Elem
instance TraverseCode Seq.Digit
instance TraverseCode Seq.Node
instance TraverseCode Seq.FingerTree
instance TraverseCode Seq.Seq where
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
instance TraverseCode Seq.ViewL
instance TraverseCode Seq.ViewR
instance TraverseCode ZipList
instance TraverseCode Complex.Complex
instance TraverseCode S.First
instance TraverseCode M.First
instance TraverseCode S.Last
instance TraverseCode M.Last
instance TraverseCode S.Min
instance TraverseCode S.Max
instance TraverseCode Ord.Down
instance TraverseCode S.WrappedMonoid
instance TraverseCode S.Dual
instance TraverseCode S.Product
instance TraverseCode S.Sum
instance TraverseCode Solo
instance TraverseCode Tree.Tree
instance TraverseCode NonEmpty.NonEmpty
instance TraverseCode m => TraverseCode (WrappedMonad m)
instance TraverseCode (p a) => TraverseCode (WrappedArrow p a)
instance TH.Lift a => TraverseCode (S.Arg a)
instance TraverseCode Proxy
instance TraverseCode f => TraverseCode (M.Ap f)
instance TraverseCode f => TraverseCode (M.Alt f)
instance TraverseCode (URec Char)
instance TraverseCode (URec Double)
instance TraverseCode (URec Float)
instance TraverseCode (URec Int)
instance TraverseCode (URec Word)
instance (Ar.Ix i, TH.Lift i) => TraverseCode (Ar.Array i) where
traverseCode :: forall (m :: * -> *) a b.
Quote m =>
(a -> Code m b) -> Array i a -> Code m (Array i b)
traverseCode a -> Code m b
f Array i a
xs = [|| Ar.listArray bnds $$csc ||]
where
csc :: Code m [b]
csc = (a -> Code m b) -> [a] -> Code m [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 (Array i a -> [a]
forall i e. Array i e -> [e]
Ar.elems Array i a
xs)
bnds :: (i, i)
bnds = Array i a -> (i, i)
forall i e. Array i e -> (i, i)
Ar.bounds Array i a
xs
instance TraverseCode PAr.Array where
traverseCode :: forall (m :: * -> *) a b.
Quote m =>
(a -> Code m b) -> Array a -> Code m (Array b)
traverseCode a -> Code m b
f Array a
ary = case [a]
lst of
[] -> [|| PAr.emptyArray ||]
[a
x] -> [|| pure $$(f x) ||]
a
x : [a]
xs -> [|| unsafeArrayFromListN len $$(f x) $$(traverseCode f xs) ||]
where
len :: Int
len = Array a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Array a
ary
lst :: [a]
lst = Array a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList Array a
ary
instance TraverseCode PSmAr.SmallArray where
traverseCode :: forall (m :: * -> *) a b.
Quote m =>
(a -> Code m b) -> SmallArray a -> Code m (SmallArray b)
traverseCode a -> Code m b
f SmallArray a
ary = case [a]
lst of
[] -> [|| PSmAr.emptySmallArray ||]
[a
x] -> [|| pure $$(f x) ||]
a
x : [a]
xs -> [|| unsafeSmallArrayFromListN len $$(f x) $$(traverseCode f xs) ||]
where
len :: Int
len = SmallArray a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length SmallArray a
ary
lst :: [a]
lst = SmallArray a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList SmallArray a
ary
unsafeArrayFromListN :: Int -> a -> [a] -> PAr.Array a
unsafeArrayFromListN :: forall a. Int -> a -> [a] -> Array a
unsafeArrayFromListN Int
n a
y [a]
ys =
Int -> a -> (forall s. MutableArray s a -> ST s ()) -> Array a
forall a.
Int -> a -> (forall s. MutableArray s a -> ST s ()) -> Array a
PAr.createArray Int
n a
y ((forall s. MutableArray s a -> ST s ()) -> Array a)
-> (forall s. MutableArray s a -> ST s ()) -> Array a
forall a b. (a -> b) -> a -> b
$ \MutableArray s a
ma ->
let go :: Int -> [a] -> ST s ()
go !Int
_ix [] = () -> ST s ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
go !Int
ix (a
x : [a]
xs) = do
MutableArray (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
PAr.writeArray MutableArray s a
MutableArray (PrimState (ST s)) a
ma Int
ix a
x
Int -> [a] -> ST s ()
go (Int
ixInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) [a]
xs
in Int -> [a] -> ST s ()
go Int
1 [a]
ys
unsafeSmallArrayFromListN :: Int -> a -> [a] -> PSmAr.SmallArray a
unsafeSmallArrayFromListN :: forall a. Int -> a -> [a] -> SmallArray a
unsafeSmallArrayFromListN Int
n a
y [a]
ys =
Int
-> a
-> (forall s. SmallMutableArray s a -> ST s ())
-> SmallArray a
forall a.
Int
-> a
-> (forall s. SmallMutableArray s a -> ST s ())
-> SmallArray a
PSmAr.createSmallArray Int
n a
y ((forall s. SmallMutableArray s a -> ST s ()) -> SmallArray a)
-> (forall s. SmallMutableArray s a -> ST s ()) -> SmallArray a
forall a b. (a -> b) -> a -> b
$ \SmallMutableArray s a
ma ->
let go :: Int -> [a] -> ST s ()
go !Int
_ix [] = () -> ST s ()
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
go !Int
ix (a
x : [a]
xs) = do
SmallMutableArray (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
PSmAr.writeSmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
ma Int
ix a
x
Int -> [a] -> ST s ()
go (Int
ixInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) [a]
xs
in Int -> [a] -> ST s ()
go Int
1 [a]
ys