{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}

module Hasql.Interpolate.Internal.EncodeRow.TH
  ( genEncodeRowInstance,
  )
where

import Control.Monad
import Data.Foldable (foldl')
import Data.Functor.Contravariant
import qualified Hasql.Encoders as E
import Hasql.Interpolate.Internal.Encoder (EncodeField (..))
import Language.Haskell.TH

-- | Generate a single 'Hasql.Interpolate.EncodeRow' instance for a
-- tuple of size @tupSize@
genEncodeRowInstance ::
  -- | tuple size
  Int ->
  Q Dec
genEncodeRowInstance :: Int -> Q Dec
genEncodeRowInstance Int
tupSize
  | Int
tupSize forall a. Ord a => a -> a -> Bool
< Int
2 = forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"this is just for tuples, must specify a tuple size of 2 or greater"
  | Bool
otherwise = do
      [Name]
tyVars <- forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
tupSize (forall (m :: * -> *). Quote m => String -> m Name
newName String
"x")
      [Type]
context <- forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (\Name
x -> [t|EncodeField $(varT x)|]) [Name]
tyVars
      let unzipWithEncoderName :: Name
unzipWithEncoderName = String -> Name
mkName String
"unzipWithEncoder"
      Type
instanceHead <- [t|$(conT (mkName "EncodeRow")) $(pure $ foldl' AppT (TupleT tupSize) (map VarT tyVars))|]
      Name
innerContName <- forall (m :: * -> *). Quote m => String -> m Name
newName String
"k"
      Exp
cons <- [e|(:)|]
      [Name]
kconsTailNames <- forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (\Name
_ -> forall (m :: * -> *). Quote m => String -> m Name
newName String
"tail") [Name]
tyVars
      let kconsPats :: [Pat]
          kconsPats :: [Pat]
kconsPats =
            [ [Pat] -> Pat
TupP (forall a b. (a -> b) -> [a] -> [b]
map Name -> Pat
VarP [Name]
tyVars),
              Pat -> Pat
TildeP ([Pat] -> Pat
TupP (forall a b. (a -> b) -> [a] -> [b]
map Name -> Pat
VarP [Name]
kconsTailNames))
            ]
          kconsTupBody :: [Exp]
          kconsTupBody :: [Exp]
kconsTupBody =
            let vars :: [Exp]
vars = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Name -> Name -> Exp
phi [Name]
tyVars [Name]
kconsTailNames
                phi :: Name -> Name -> Exp
phi Name
headName Name
tailName = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Exp -> Exp -> Exp
AppE Exp
cons [Name -> Exp
VarE Name
headName, Name -> Exp
VarE Name
tailName]
             in [Exp]
vars
          kcons :: Exp
          kcons :: Exp
kcons = [Pat] -> Exp -> Exp
LamE [Pat]
kconsPats ([Maybe Exp] -> Exp
TupE (forall a b. (a -> b) -> [a] -> [b]
map forall a. a -> Maybe a
Just [Exp]
kconsTupBody))
          knil :: Exp
          knil :: Exp
knil = [Maybe Exp] -> Exp
TupE forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a. Int -> a -> [a]
replicate Int
tupSize ([Exp] -> Exp
ListE [])
      Exp
kenc :: Exp <- do
        let listEncoder :: Q Exp
listEncoder = [e|E.param (E.nonNullable (E.foldableArray encodeField))|]
            plucks :: [Q Exp]
plucks = forall a b. (a -> b) -> [a] -> [b]
map (Int -> Int -> Q Exp
pluck Int
tupSize) [Int
0 .. Int
tupSize forall a. Num a => a -> a -> a
- Int
1]
        [Exp]
encExps <- forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (\Q Exp
getTupElem -> [e|contramap $getTupElem $listEncoder|]) [Q Exp]
plucks
        forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Exp
a Q Exp
b -> [e|$(pure a) <> $(b)|]) [e|mempty|] [Exp]
encExps
      let kExp :: Exp
          kExp :: Exp
kExp = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Exp -> Exp -> Exp
AppE (Name -> Exp
VarE Name
innerContName) [Exp
kcons, Exp
knil, Exp
kenc, Lit -> Exp
LitE (Integer -> Lit
IntegerL (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
tupSize))]
      let instanceBody :: Dec
instanceBody = Name -> [Clause] -> Dec
FunD Name
unzipWithEncoderName [[Pat] -> Body -> [Dec] -> Clause
Clause [Name -> Pat
VarP Name
innerContName] (Exp -> Body
NormalB Exp
kExp) []]
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe Overlap -> [Type] -> Type -> [Dec] -> Dec
InstanceD forall a. Maybe a
Nothing [Type]
context Type
instanceHead [Dec
instanceBody])

pluck :: Int -> Int -> Q Exp
pluck :: Int -> Int -> Q Exp
pluck Int
1 Int
0 = [e|id|]
pluck Int
tupSize Int
idx = do
  Name
matchName <- forall (m :: * -> *). Quote m => String -> m Name
newName String
"match"
  let tupPat :: Pat
tupPat = [Pat] -> Pat
TupP (forall a b. (a -> b) -> [a] -> [b]
map (\Int
n -> if Int
n forall a. Eq a => a -> a -> Bool
== Int
idx then Name -> Pat
VarP Name
matchName else Pat
WildP) [Int
0 .. Int
tupSize forall a. Num a => a -> a -> a
- Int
1])
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [Pat] -> Exp -> Exp
LamE [Pat
tupPat] (Name -> Exp
VarE Name
matchName)