{-# 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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
2 = String -> Q Dec
forall a. String -> Q a
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
      tyVars <- Int -> Q Name -> Q [Name]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
tupSize (String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"x")
      context <- traverse (\Name
x -> [t|EncodeField $(Name -> Q Type
forall (m :: * -> *). Quote m => Name -> m Type
varT Name
x)|]) tyVars
      let unzipWithEncoderName = String -> Name
mkName String
"unzipWithEncoder"
      instanceHead <- [t|$(conT (mkName "EncodeRow")) $(pure $ foldl' AppT (TupleT tupSize) (map VarT tyVars))|]
      innerContName <- newName "k"
      cons <- [e|(:)|]
      kconsTailNames <- traverse (\Name
_ -> String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"tail") tyVars
      let 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 =
            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 b a. (b -> a -> b) -> b -> [a] -> b
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 = [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 = [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 [])
      kenc :: Exp <- do
        let listEncoder = [e|E.param (E.nonNullable (E.foldableArray encodeField))|]
            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]
        encExps <- traverse (\Q Exp
getTupElem -> [e|contramap $Q Exp
getTupElem $Q Exp
listEncoder|]) plucks
        foldr (\Exp
a Q Exp
b -> [e|$(Exp -> Q Exp
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
a) <> $(Q Exp
b)|]) [e|mempty|] encExps
      let kExp :: Exp
          kExp = (Exp -> Exp -> Exp) -> Exp -> [Exp] -> Exp
forall b a. (b -> a -> b) -> b -> [a] -> b
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 = Name -> [Clause] -> Dec
FunD Name
unzipWithEncoderName [[Pat] -> Body -> [Dec] -> Clause
Clause [Name -> Pat
VarP Name
innerContName] (Exp -> Body
NormalB Exp
kExp) []]
      pure (InstanceD Nothing context instanceHead [instanceBody])

pluck :: Int -> Int -> Q Exp
pluck :: Int -> Int -> Q Exp
pluck Int
1 Int
0 = [e|id|]
pluck Int
tupSize Int
idx = do
  matchName <- String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"match"
  let 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])
  pure $ LamE [tupPat] (VarE matchName)