{-# 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 (..), EncodeValue (..))
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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2 = String -> Q Dec
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 <- Int -> Q Name -> Q [Name]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
tupSize (String -> Q Name
newName String
"x")
    [Type]
context <- (Name -> Q Type) -> [Name] -> Q [Type]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (\Name
x -> [t|EncodeValue $(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 <- String -> Q Name
newName String
"k"
    Exp
cons <- [e|(:)|]
    [Name]
kconsTailNames <- (Name -> Q Name) -> [Name] -> Q [Name]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (\Name
_ -> String -> Q Name
newName String
"tail") [Name]
tyVars
    let kconsPats :: [Pat]
        kconsPats :: [Pat]
kconsPats =
          [ [Pat] -> Pat
TupP ((Name -> Pat) -> [Name] -> [Pat]
forall a b. (a -> b) -> [a] -> [b]
map Name -> Pat
VarP [Name]
tyVars),
            Pat -> Pat
TildeP ([Pat] -> Pat
TupP ((Name -> Pat) -> [Name] -> [Pat]
forall a b. (a -> b) -> [a] -> [b]
map Name -> Pat
VarP [Name]
kconsTailNames))
          ]
        kconsTupBody :: [Exp]
        kconsTupBody :: [Exp]
kconsTupBody =
          let vars :: [Exp]
vars = (Name -> Name -> Exp) -> [Name] -> [Name] -> [Exp]
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 = (Exp -> Exp -> Exp) -> Exp -> [Exp] -> Exp
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 ((Exp -> Maybe Exp) -> [Exp] -> [Maybe Exp]
forall a b. (a -> b) -> [a] -> [b]
map Exp -> Maybe Exp
forall a. a -> Maybe a
Just [Exp]
kconsTupBody))
        knil :: Exp
        knil :: Exp
knil = [Maybe Exp] -> Exp
TupE ([Maybe Exp] -> Exp) -> ([Exp] -> [Maybe Exp]) -> [Exp] -> Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Exp -> Maybe Exp) -> [Exp] -> [Maybe Exp]
forall a b. (a -> b) -> [a] -> [b]
map Exp -> Maybe Exp
forall a. a -> Maybe a
Just ([Exp] -> Exp) -> [Exp] -> Exp
forall a b. (a -> b) -> a -> b
$ Int -> Exp -> [Exp]
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 = (Int -> Q Exp) -> [Int] -> [Q Exp]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Int -> Q Exp
pluck Int
tupSize) [Int
0 .. Int
tupSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
      [Exp]
encExps <- (Q Exp -> Q Exp) -> [Q Exp] -> Q [Exp]
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
      (Exp -> Q Exp -> Q Exp) -> Q Exp -> [Exp] -> Q Exp
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 = (Exp -> Exp -> Exp) -> Exp -> [Exp] -> Exp
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 (Int -> Integer
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) []]
    Dec -> Q Dec
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe Overlap -> [Type] -> Type -> [Dec] -> Dec
InstanceD Maybe Overlap
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 <- String -> Q Name
newName String
"match"
  let tupPat :: Pat
tupPat = [Pat] -> Pat
TupP ((Int -> Pat) -> [Int] -> [Pat]
forall a b. (a -> b) -> [a] -> [b]
map (\Int
n -> if Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
idx then Name -> Pat
VarP Name
matchName else Pat
WildP) [Int
0 .. Int
tupSize Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1])
  Exp -> Q Exp
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> Q Exp) -> Exp -> Q Exp
forall a b. (a -> b) -> a -> b
$ [Pat] -> Exp -> Exp
LamE [Pat
tupPat] (Name -> Exp
VarE Name
matchName)