-- | Helpers for generating instances of the 'Record' type class.
--
-- @since 0.0.1.0
module Prairie.TH where

import Data.Constraint (Dict(..))
import Language.Haskell.TH
import Control.Lens (lens)
import qualified Data.List as List
import Data.Traversable (for)
import Data.Char (toUpper, toLower)
import qualified Data.Text as Text

import Prairie.Class

-- | Create an instance of the 'Record' type class.
--
-- @
-- data User
--   = User
--   { name :: String
--   , age :: Int
--   }
--
-- mkRecord ''User
--
-- ====>
--
-- instance Record User where
--   data Field User a where
--     UserName :: String
--     UserAge :: Int
--
--   recordFieldLens fl =
--     case fl of
--       UserName -> lens name (\u n -> u { name = n)
--       UserAge -> lens age (\u n -> u { age = n)
--
-- instance SymbolToField "age" User Int where symbolToField = UserName
-- instance SymbolToField "name" User String where symbolToField = UserAge
-- @
--
-- If the fields are prefixed with the type's name, this function figures
-- it out and won't duplicate the field.
--
-- @
-- data User
--   = User
--   { userName :: String
--   , userAge :: Int
--   }
--
-- mkRecord ''User
--
-- ====>
--
-- instance Record User where
--   data Field User a where
--     UserName :: String
--     UserAge :: Int
--
--   recordFieldLens fl =
--     case fl of
--       UserName -> lens name (\u n -> u { name = n)
--       UserAge -> lens age (\u n -> u { age = n)
--
-- instance SymbolToField "name" User Int where symbolToField = UserName
-- instance SymbolToField "age" User String where symbolToField = UserAge
-- @
--
-- @since 0.0.1.0
mkRecord :: Name -> DecsQ
mkRecord :: Name -> DecsQ
mkRecord Name
u = do
  Info
ty <- Name -> Q Info
reify Name
u
  (Name
typeName, Con
con) <-
    case Info
ty of
      TyConI Dec
dec ->
        case Dec
dec of
          DataD Cxt
_cxt Name
name [TyVarBndr]
_tyvars Maybe Kind
_mkind [Con
con] [DerivClause]
_derivs ->
            (Name, Con) -> Q (Name, Con)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Name
name, Con
con)
          NewtypeD Cxt
_cxt Name
name [TyVarBndr]
_tyvars Maybe Kind
_mkind Con
con [DerivClause]
_derivs ->
            (Name, Con) -> Q (Name, Con)
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Name
name, Con
con)
          Dec
_ ->
            String -> Q (Name, Con)
forall (m :: Type -> Type) a. MonadFail m => String -> m a
fail String
"unsupported data structure"
      Info
_ ->
        String -> Q (Name, Con)
forall (m :: Type -> Type) a. MonadFail m => String -> m a
fail String
"unsupported type"

  let
    stripTypeName :: Name -> Name
stripTypeName Name
n =
      let
        typeNamePrefix :: String
typeNamePrefix =
          String -> String
lowerFirst (Name -> String
nameBase Name
typeName)
       in
        case String -> String -> Maybe String
forall a. Eq a => [a] -> [a] -> Maybe [a]
List.stripPrefix String
typeNamePrefix (Name -> String
nameBase Name
n) of
          Just String
xs -> String -> Name
mkName (String -> String
lowerFirst String
xs)
          Maybe String
Nothing -> Name
n

  (Name
recordCon, [(Name, Kind)]
names'types) <-
    case Con
con of
      RecC Name
conName [VarBangType]
varBangTypes ->
        (Name, [(Name, Kind)]) -> Q (Name, [(Name, Kind)])
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure ((Name, [(Name, Kind)]) -> Q (Name, [(Name, Kind)]))
-> (Name, [(Name, Kind)]) -> Q (Name, [(Name, Kind)])
forall a b. (a -> b) -> a -> b
$ (Name
conName, (VarBangType -> (Name, Kind)) -> [VarBangType] -> [(Name, Kind)]
forall a b. (a -> b) -> [a] -> [b]
map (\(Name
n, Bang
_b, Kind
t) -> (Name
n, Kind
t)) [VarBangType]
varBangTypes)
      Con
_ ->
        String -> Q (Name, [(Name, Kind)])
forall (m :: Type -> Type) a. MonadFail m => String -> m a
fail String
"only supports records"

  let
    mkConstrFieldName :: Name -> Name
mkConstrFieldName Name
fieldName =
      String -> Name
mkName (Name -> String
nameBase Name
typeName String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String -> String
upperFirst (Name -> String
nameBase (Name -> Name
stripTypeName Name
fieldName)))

  Clause
fieldLensClause <- do
    Name
arg <- String -> Q Name
newName String
"field"
    let
      mkMatch :: (Name, Kind) -> Q Match
mkMatch (Name
fieldName, Kind
_typ) = do
        Name
recVar <- String -> Q Name
newName String
"rec"
        Name
newVal <- String -> Q Name
newName String
"newVal"

        Match -> Q Match
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Match -> Q Match) -> Match -> Q Match
forall a b. (a -> b) -> a -> b
$
          Pat -> Body -> [Dec] -> Match
Match
            (Name -> [Pat] -> Pat
ConP (Name -> Name
mkConstrFieldName Name
fieldName) [])
            (Exp -> Body
NormalB (Exp -> Body) -> Exp -> Body
forall a b. (a -> b) -> a -> b
$
            Name -> Exp
VarE 'lens
            Exp -> Exp -> Exp
`AppE` Name -> Exp
VarE Name
fieldName
            Exp -> Exp -> Exp
`AppE`
              [Pat] -> Exp -> Exp
LamE [Name -> Pat
VarP Name
recVar, Name -> Pat
VarP Name
newVal]
                (Exp -> [FieldExp] -> Exp
RecUpdE (Name -> Exp
VarE Name
recVar) [(Name
fieldName, Name -> Exp
VarE Name
newVal)])
            )
            []
    Exp
body <- Exp -> [Match] -> Exp
CaseE (Name -> Exp
VarE Name
arg) ([Match] -> Exp) -> Q [Match] -> Q Exp
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Name, Kind) -> Q Match) -> [(Name, Kind)] -> Q [Match]
forall (t :: Type -> Type) (f :: Type -> Type) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (Name, Kind) -> Q Match
mkMatch [(Name, Kind)]
names'types
    Clause -> Q Clause
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Clause -> Q Clause) -> Clause -> Q Clause
forall a b. (a -> b) -> a -> b
$ [Pat] -> Body -> [Dec] -> Clause
Clause [Name -> Pat
VarP Name
arg] (Exp -> Body
NormalB Exp
body) []
  let
    recordFieldLensDec :: Dec
recordFieldLensDec =
      Name -> [Clause] -> Dec
FunD 'recordFieldLens [Clause
fieldLensClause]
    fieldConstructors :: [(Name, Kind)]
fieldConstructors =
      ((Name, Kind) -> (Name, Kind)) -> [(Name, Kind)] -> [(Name, Kind)]
forall a b. (a -> b) -> [a] -> [b]
map (\(Name
n, Kind
t) -> (Name -> Name
mkConstrFieldName Name
n, Kind
t)) [(Name, Kind)]
names'types

  Dec
mkAllFields <- Dec -> Q Dec
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Dec -> Q Dec) -> Dec -> Q Dec
forall a b. (a -> b) -> a -> b
$
    Pat -> Body -> [Dec] -> Dec
ValD
      (Name -> Pat
VarP 'allFields)
      (Exp -> Body
NormalB (Exp -> Body) -> Exp -> Body
forall a b. (a -> b) -> a -> b
$ [Exp] -> Exp
ListE (((Name, Kind) -> Exp) -> [(Name, Kind)] -> [Exp]
forall a b. (a -> b) -> [a] -> [b]
map (Exp -> Exp -> Exp
AppE (Name -> Exp
ConE 'SomeField) (Exp -> Exp) -> ((Name, Kind) -> Exp) -> (Name, Kind) -> Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> Exp
ConE (Name -> Exp) -> ((Name, Kind) -> Name) -> (Name, Kind) -> Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, Kind) -> Name
forall a b. (a, b) -> a
fst) [(Name, Kind)]
fieldConstructors))
      []

  Dec
mkTabulateRecord <- do
    Name
fromFieldName <- String -> Q Name
newName String
"fromField"
    Exp
body <- Exp -> Q Exp
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Exp -> Q Exp) -> Exp -> Q Exp
forall a b. (a -> b) -> a -> b
$
      Name -> [FieldExp] -> Exp
RecConE Name
recordCon ([FieldExp] -> Exp) -> [FieldExp] -> Exp
forall a b. (a -> b) -> a -> b
$
        ((Name, Kind) -> FieldExp) -> [(Name, Kind)] -> [FieldExp]
forall a b. (a -> b) -> [a] -> [b]
map
          (\(Name
n, Kind
_) -> (Name
n, Name -> Exp
VarE Name
fromFieldName Exp -> Exp -> Exp
`AppE` Name -> Exp
ConE (Name -> Name
mkConstrFieldName Name
n)))
          [(Name, Kind)]
names'types

    Dec -> Q Dec
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Dec -> Q Dec) -> Dec -> Q Dec
forall a b. (a -> b) -> a -> b
$
      Name -> [Clause] -> Dec
FunD 'tabulateRecord
        [ [Pat] -> Body -> [Dec] -> Clause
Clause [Name -> Pat
VarP Name
fromFieldName] (Exp -> Body
NormalB Exp
body) []
        ]

  Dec
mkRecordFieldLabel <- do
    Name
fieldName <- String -> Q Name
newName String
"fieldName"
    Exp
body <- Exp -> Q Exp
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Exp -> Q Exp) -> Exp -> Q Exp
forall a b. (a -> b) -> a -> b
$
      Exp -> [Match] -> Exp
CaseE (Name -> Exp
VarE Name
fieldName)  ([Match] -> Exp) -> [Match] -> Exp
forall a b. (a -> b) -> a -> b
$
        (((Name, Kind) -> Match) -> [(Name, Kind)] -> [Match])
-> [(Name, Kind)] -> ((Name, Kind) -> Match) -> [Match]
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Name, Kind) -> Match) -> [(Name, Kind)] -> [Match]
forall a b. (a -> b) -> [a] -> [b]
map [(Name, Kind)]
names'types (((Name, Kind) -> Match) -> [Match])
-> ((Name, Kind) -> Match) -> [Match]
forall a b. (a -> b) -> a -> b
$ \(Name
n, Kind
_) ->
          let
            constrFieldName :: Name
constrFieldName =
              Name -> Name
mkConstrFieldName Name
n
            pat :: Pat
pat =
              Name -> [Pat] -> Pat
ConP Name
constrFieldName []
            bdy :: Exp
bdy =
              Exp -> Exp -> Exp
AppE (Name -> Exp
VarE 'Text.pack) (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ Lit -> Exp
LitE (Lit -> Exp) -> Lit -> Exp
forall a b. (a -> b) -> a -> b
$ String -> Lit
StringL (String -> Lit) -> String -> Lit
forall a b. (a -> b) -> a -> b
$ Name -> String
nameBase (Name -> String) -> Name -> String
forall a b. (a -> b) -> a -> b
$ Name -> Name
stripTypeName Name
n

           in
            Pat -> Body -> [Dec] -> Match
Match Pat
pat (Exp -> Body
NormalB Exp
bdy)  []
    Dec -> Q Dec
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Dec -> Q Dec) -> Dec -> Q Dec
forall a b. (a -> b) -> a -> b
$
      Name -> [Clause] -> Dec
FunD 'recordFieldLabel
        [ [Pat] -> Body -> [Dec] -> Clause
Clause [Name -> Pat
VarP Name
fieldName] (Exp -> Body
NormalB Exp
body) []
        ]

  let
    fieldConstrs :: [Con]
fieldConstrs =
      ((Name, Kind) -> Con) -> [(Name, Kind)] -> [Con]
forall a b. (a -> b) -> [a] -> [b]
map (Name, Kind) -> Con
mkFieldConstr [(Name, Kind)]
fieldConstructors
    mkFieldConstr :: (Name, Kind) -> Con
mkFieldConstr (Name
fieldName, Kind
typ) =
      [Name] -> [BangType] -> Kind -> Con
GadtC
        [ Name
fieldName
        ]
        []
        (Name -> Kind
ConT ''Field Kind -> Kind -> Kind
`AppT` Name -> Kind
ConT Name
typeName Kind -> Kind -> Kind
`AppT` Kind
typ)

    recordInstance :: Dec
recordInstance =
      Maybe Overlap -> Cxt -> Kind -> [Dec] -> Dec
InstanceD
        Maybe Overlap
forall a. Maybe a
Nothing
        []
        (Name -> Kind
ConT ''Record Kind -> Kind -> Kind
`AppT` Name -> Kind
ConT Name
typeName)
        (
          [ Cxt
-> Maybe [TyVarBndr]
-> Kind
-> Maybe Kind
-> [Con]
-> [DerivClause]
-> Dec
DataInstD
              []
              Maybe [TyVarBndr]
forall a. Maybe a
Nothing
              (Name -> Kind
ConT ''Field Kind -> Kind -> Kind
`AppT` Name -> Kind
ConT Name
typeName Kind -> Kind -> Kind
`AppT` Name -> Kind
VarT (String -> Name
mkName String
"a"))
              Maybe Kind
forall a. Maybe a
Nothing
              [Con]
fieldConstrs
              []
          , Dec
recordFieldLensDec
          , Dec
mkAllFields
          , Dec
mkTabulateRecord
          , Dec
mkRecordFieldLabel
          ]
        )

  Dec
fieldDictInstance <- do
    Name
constraintVar <- String -> Q Name
newName String
"c"
    Name
fieldVar <- String -> Q Name
newName String
"field"
    let
      allFieldsC :: Cxt
allFieldsC =
        (Kind -> Kind) -> Cxt -> Cxt
forall a b. (a -> b) -> [a] -> [b]
map (Name -> Kind
VarT Name
constraintVar Kind -> Kind -> Kind
`AppT`) (((Name, Kind) -> Kind) -> [(Name, Kind)] -> Cxt
forall a b. (a -> b) -> [a] -> [b]
map (Name, Kind) -> Kind
forall a b. (a, b) -> b
snd [(Name, Kind)]
names'types)
      fieldDictDecl :: [Dec]
fieldDictDecl =
        [ Name -> [Clause] -> Dec
FunD 'getFieldDict [[Pat] -> Body -> [Dec] -> Clause
Clause [Name -> Pat
VarP Name
fieldVar] (Exp -> Body
NormalB Exp
fieldDictBody) []]
        ]
      fieldDictBody :: Exp
fieldDictBody =
        Exp -> [Match] -> Exp
CaseE (Name -> Exp
VarE Name
fieldVar) ([Match] -> Exp) -> [Match] -> Exp
forall a b. (a -> b) -> a -> b
$ ((Name, Kind) -> Match) -> [(Name, Kind)] -> [Match]
forall a b. (a -> b) -> [a] -> [b]
map (Name, Kind) -> Match
forall b. (Name, b) -> Match
mkFieldDictMatches [(Name, Kind)]
fieldConstructors
      mkFieldDictMatches :: (Name, b) -> Match
mkFieldDictMatches (Name
name, b
_type) =
        Pat -> Body -> [Dec] -> Match
Match (Name -> [Pat] -> Pat
ConP Name
name []) (Exp -> Body
NormalB (Name -> Exp
ConE 'Dict)) []

    Dec -> Q Dec
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Dec -> Q Dec) -> Dec -> Q Dec
forall a b. (a -> b) -> a -> b
$
      Maybe Overlap -> Cxt -> Kind -> [Dec] -> Dec
InstanceD
        Maybe Overlap
forall a. Maybe a
Nothing -- maybe overlap
        Cxt
allFieldsC
        (Name -> Kind
ConT ''FieldDict Kind -> Kind -> Kind
`AppT` Name -> Kind
VarT Name
constraintVar Kind -> Kind -> Kind
`AppT` Name -> Kind
ConT Name
typeName)
        [Dec]
fieldDictDecl

  [Dec]
symbolToFieldInstances <-
    ([[Dec]] -> [Dec]) -> Q [[Dec]] -> DecsQ
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap [[Dec]] -> [Dec]
forall (t :: Type -> Type) a. Foldable t => t [a] -> [a]
concat (Q [[Dec]] -> DecsQ) -> Q [[Dec]] -> DecsQ
forall a b. (a -> b) -> a -> b
$ [(Name, Kind)] -> ((Name, Kind) -> DecsQ) -> Q [[Dec]]
forall (t :: Type -> Type) (f :: Type -> Type) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for [(Name, Kind)]
names'types (((Name, Kind) -> DecsQ) -> Q [[Dec]])
-> ((Name, Kind) -> DecsQ) -> Q [[Dec]]
forall a b. (a -> b) -> a -> b
$ \(Name
fieldName, Kind
typ) -> do
      [d|
        instance SymbolToField $(litT (strTyLit (nameBase fieldName))) $(conT typeName) $(pure typ) where
          symbolToField = $(conE (mkConstrFieldName fieldName))
        |]

  [Dec] -> DecsQ
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure ([Dec] -> DecsQ) -> [Dec] -> DecsQ
forall a b. (a -> b) -> a -> b
$
    [ Dec
recordInstance
    , Dec
fieldDictInstance
    ]
    [Dec] -> [Dec] -> [Dec]
forall a. [a] -> [a] -> [a]
++
      [Dec]
symbolToFieldInstances

overFirst :: (Char -> Char) -> String -> String
overFirst :: (Char -> Char) -> String -> String
overFirst Char -> Char
f String
str =
  case String
str of
    [] -> []
    (Char
c:String
cs) -> Char -> Char
f Char
c Char -> String -> String
forall a. a -> [a] -> [a]
: String
cs

upperFirst, lowerFirst :: String -> String
upperFirst :: String -> String
upperFirst = (Char -> Char) -> String -> String
overFirst Char -> Char
toUpper
lowerFirst :: String -> String
lowerFirst = (Char -> Char) -> String -> String
overFirst Char -> Char
toLower