{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskellQuotes #-}
{-# LANGUAGE TypeApplications #-}

{- |
Template-Haskell @deriveSchema@ splice for "DataFrame.Internal.Schema".
Kept in a separate module so the runtime schema types in
"DataFrame.Internal.Schema" do not pull in @template-haskell@.
-}
module DataFrame.Internal.Schema.TH (
    deriveSchema,
    camelToSnake,
) where

import Data.Char (isUpper, toLower, toUpper)
import qualified Data.Text as T
import Language.Haskell.TH

import DataFrame.Internal.Expression (Expr)
import DataFrame.Operators (col)
import DataFrame.Schema (Schema, makeSchema, schemaType)

{- | Auto-generate a runtime 'Schema' (and per-column @'Expr'@ accessors)
from a record ADT. Emits @\<tyName\>Schema@ plus one accessor per field
(@camelCase -> snake_case@). Requires a single record constructor.
-}
deriveSchema :: Name -> DecsQ
deriveSchema :: Name -> DecsQ
deriveSchema Name
tyName = do
    Info
info <- Name -> Q Info
reify Name
tyName
    [VarBangType]
fields <- Name -> Info -> Q [VarBangType]
extractRecordFields Name
tyName Info
info
    let entries :: [([Char], [Char], Type)]
entries =
            [ ([Char] -> [Char]
camelToSnake [Char]
fieldBase, [Char]
fieldBase, Type
fTy)
            | (Name
fName, Bang
_bang, Type
fTy) <- [VarBangType]
fields
            , let fieldBase :: [Char]
fieldBase = Name -> [Char]
nameBase Name
fName
            ]
        schemaName :: Name
schemaName = [Char] -> Name
mkName ([Char] -> [Char]
lowerFirst (Name -> [Char]
nameBase Name
tyName) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"Schema")
        prefix :: [Char]
prefix = [Char] -> [Char]
lowerFirst (Name -> [Char]
nameBase Name
tyName)
        tupleE :: ([Char], b, Type) -> Exp
tupleE ([Char]
colName, b
_, Type
fTy) =
            [Maybe Exp] -> Exp
TupE
                [ Exp -> Maybe Exp
forall a. a -> Maybe a
Just (Exp -> Exp -> Exp
AppE (Name -> Exp
VarE 'T.pack) (Lit -> Exp
LitE ([Char] -> Lit
StringL [Char]
colName)))
                , Exp -> Maybe Exp
forall a. a -> Maybe a
Just (Exp -> Type -> Exp
AppTypeE (Name -> Exp
VarE 'schemaType) Type
fTy)
                ]
        schemaBody :: Exp
schemaBody =
            Exp -> Exp -> Exp
AppE (Name -> Exp
VarE 'makeSchema) ([Exp] -> Exp
ListE ((([Char], [Char], Type) -> Exp)
-> [([Char], [Char], Type)] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map ([Char], [Char], Type) -> Exp
forall {b}. ([Char], b, Type) -> Exp
tupleE [([Char], [Char], Type)]
entries))
        schemaDecls :: [Dec]
schemaDecls =
            [ Name -> Type -> Dec
SigD Name
schemaName (Name -> Type
ConT ''Schema)
            , Pat -> Body -> [Dec] -> Dec
ValD (Name -> Pat
VarP Name
schemaName) (Exp -> Body
NormalB Exp
schemaBody) []
            ]
        accessorDecls :: [Dec]
accessorDecls =
            [[Dec]] -> [Dec]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
                [ [ Name -> Type -> Dec
SigD Name
accName (Type -> Type -> Type
AppT (Name -> Type
ConT ''Expr) Type
fTy)
                  , Pat -> Body -> [Dec] -> Dec
ValD
                        (Name -> Pat
VarP Name
accName)
                        ( Exp -> Body
NormalB
                            ( Exp -> Exp -> Exp
AppE
                                (Name -> Exp
VarE 'col)
                                ( Exp -> Exp -> Exp
AppE
                                    (Name -> Exp
VarE 'T.pack)
                                    (Lit -> Exp
LitE ([Char] -> Lit
StringL [Char]
colName))
                                )
                            )
                        )
                        []
                  ]
                | ([Char]
colName, [Char]
fieldBase, Type
fTy) <- [([Char], [Char], Type)]
entries
                , let accName :: Name
accName = [Char] -> Name
mkName ([Char]
prefix [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char] -> [Char]
upperFirst [Char]
fieldBase)
                ]
    [Dec] -> DecsQ
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Dec]
schemaDecls [Dec] -> [Dec] -> [Dec]
forall a. [a] -> [a] -> [a]
++ [Dec]
accessorDecls)

extractRecordFields :: Name -> Info -> Q [VarBangType]
extractRecordFields :: Name -> Info -> Q [VarBangType]
extractRecordFields Name
_ (TyConI Dec
dec) = case Dec
dec of
    DataD Cxt
_ Name
_ [TyVarBndr BndrVis]
_ Maybe Type
_ [RecC Name
_ [VarBangType]
fs] [DerivClause]
_ -> [VarBangType] -> Q [VarBangType]
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [VarBangType]
fs
    NewtypeD Cxt
_ Name
_ [TyVarBndr BndrVis]
_ Maybe Type
_ (RecC Name
_ [VarBangType]
fs) [DerivClause]
_ -> [VarBangType] -> Q [VarBangType]
forall a. a -> Q a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [VarBangType]
fs
    DataD Cxt
_ Name
n [TyVarBndr BndrVis]
_ Maybe Type
_ [Con]
_ [DerivClause]
_ ->
        [Char] -> Q [VarBangType]
forall a. [Char] -> Q a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail ([Char] -> Q [VarBangType]) -> [Char] -> Q [VarBangType]
forall a b. (a -> b) -> a -> b
$
            [Char]
"deriveSchema: "
                [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Name -> [Char]
forall a. Show a => a -> [Char]
show Name
n
                [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" must have exactly one record constructor"
    NewtypeD Cxt
_ Name
n [TyVarBndr BndrVis]
_ Maybe Type
_ Con
_ [DerivClause]
_ ->
        [Char] -> Q [VarBangType]
forall a. [Char] -> Q a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail ([Char] -> Q [VarBangType]) -> [Char] -> Q [VarBangType]
forall a b. (a -> b) -> a -> b
$
            [Char]
"deriveSchema: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Name -> [Char]
forall a. Show a => a -> [Char]
show Name
n [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" newtype must use record syntax"
    Dec
other ->
        [Char] -> Q [VarBangType]
forall a. [Char] -> Q a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail ([Char] -> Q [VarBangType]) -> [Char] -> Q [VarBangType]
forall a b. (a -> b) -> a -> b
$
            [Char]
"deriveSchema: unsupported declaration: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Dec -> [Char]
forall a. Show a => a -> [Char]
show Dec
other
extractRecordFields Name
tyName Info
_ =
    [Char] -> Q [VarBangType]
forall a. [Char] -> Q a
forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail ([Char] -> Q [VarBangType]) -> [Char] -> Q [VarBangType]
forall a b. (a -> b) -> a -> b
$
        [Char]
"deriveSchema: "
            [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Name -> [Char]
forall a. Show a => a -> [Char]
show Name
tyName
            [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
" is not a data/newtype declaration"

{- | @camelCase -> snake_case@. Lowercases the first character then prefixes
@\'_\'@ before every uppercase character (lowercased).
-}
camelToSnake :: String -> String
camelToSnake :: [Char] -> [Char]
camelToSnake [] = []
camelToSnake (Char
c : [Char]
cs) = Char -> Char
toLower Char
c Char -> [Char] -> [Char]
forall a. a -> [a] -> [a]
: [Char] -> [Char]
go [Char]
cs
  where
    go :: [Char] -> [Char]
go [] = []
    go (Char
x : [Char]
xs)
        | Char -> Bool
isUpper Char
x = Char
'_' Char -> [Char] -> [Char]
forall a. a -> [a] -> [a]
: Char -> Char
toLower Char
x Char -> [Char] -> [Char]
forall a. a -> [a] -> [a]
: [Char] -> [Char]
go [Char]
xs
        | Bool
otherwise = Char
x Char -> [Char] -> [Char]
forall a. a -> [a] -> [a]
: [Char] -> [Char]
go [Char]
xs

lowerFirst :: String -> String
lowerFirst :: [Char] -> [Char]
lowerFirst [] = []
lowerFirst (Char
c : [Char]
cs) = Char -> Char
toLower Char
c Char -> [Char] -> [Char]
forall a. a -> [a] -> [a]
: [Char]
cs

upperFirst :: String -> String
upperFirst :: [Char] -> [Char]
upperFirst [] = []
upperFirst (Char
c : [Char]
cs) = Char -> Char
toUpper Char
c Char -> [Char] -> [Char]
forall a. a -> [a] -> [a]
: [Char]
cs