{-# LANGUAGE CPP #-}

-- | 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 ->
            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 ->
            forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Name
name, Con
con)
          Dec
_ ->
            forall (m :: Type -> Type) a. MonadFail m => String -> m a
fail String
"unsupported data structure"
      Info
_ ->
        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 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 ->
        forall (f :: Type -> Type) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ (Name
conName, forall a b. (a -> b) -> [a] -> [b]
map (\(Name
n, Bang
_b, Kind
t) -> (Name
n, Kind
t)) [VarBangType]
varBangTypes)
      Con
_ ->
        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 forall a. Semigroup a => a -> a -> a
<> String -> String
upperFirst (Name -> String
nameBase (Name -> Name
stripTypeName Name
fieldName)))

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

        forall (f :: Type -> Type) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
          Pat -> Body -> [Dec] -> Match
Match
            (Name -> Pat
compatConP (Name -> Name
mkConstrFieldName Name
fieldName))
            (Exp -> Body
NormalB 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) forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> 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
    forall (f :: Type -> Type) a. Applicative f => a -> f a
pure 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 =
      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 <- forall (f :: Type -> Type) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
    Pat -> Body -> [Dec] -> Dec
ValD
      (Name -> Pat
VarP 'allFields)
      (Exp -> Body
NormalB forall a b. (a -> b) -> a -> b
$ [Exp] -> Exp
ListE (forall a b. (a -> b) -> [a] -> [b]
map (Exp -> Exp -> Exp
AppE (Name -> Exp
ConE 'SomeField) forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> Exp
ConE forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [(Name, Kind)]
fieldConstructors))
      []

  Dec
mkTabulateRecord <- do
    Name
fromFieldName <- forall (m :: Type -> Type). Quote m => String -> m Name
newName String
"fromField"
    Exp
body <- forall (f :: Type -> Type) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
      Name -> [FieldExp] -> Exp
RecConE Name
recordCon forall a b. (a -> b) -> a -> b
$
        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

    forall (f :: Type -> Type) a. Applicative f => a -> f a
pure 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 <- forall (m :: Type -> Type). Quote m => String -> m Name
newName String
"fieldName"
    Exp
body <- forall (f :: Type -> Type) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
      Exp -> [Match] -> Exp
CaseE (Name -> Exp
VarE Name
fieldName)  forall a b. (a -> b) -> a -> b
$
        forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a b. (a -> b) -> [a] -> [b]
map [(Name, Kind)]
names'types forall a b. (a -> b) -> a -> b
$ \(Name
n, Kind
_) ->
          let
            constrFieldName :: Name
constrFieldName =
              Name -> Name
mkConstrFieldName Name
n
            pat :: Pat
pat =
                Name -> Pat
compatConP Name
constrFieldName
            bdy :: Exp
bdy =
              Exp -> Exp -> Exp
AppE (Name -> Exp
VarE 'Text.pack) forall a b. (a -> b) -> a -> b
$ Lit -> Exp
LitE forall a b. (a -> b) -> a -> b
$ String -> Lit
StringL forall a b. (a -> b) -> a -> b
$ Name -> String
nameBase 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)  []
    forall (f :: Type -> Type) a. Applicative f => a -> f a
pure 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 =
      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
        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
              []
              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"))
              forall a. Maybe a
Nothing
              [Con]
fieldConstrs
              []
          , Dec
recordFieldLensDec
          , Dec
mkAllFields
          , Dec
mkTabulateRecord
          , Dec
mkRecordFieldLabel
          ]
        )

  Dec
fieldDictInstance <- do
    Name
constraintVar <- forall (m :: Type -> Type). Quote m => String -> m Name
newName String
"c"
    Name
fieldVar <- forall (m :: Type -> Type). Quote m => String -> m Name
newName String
"field"
    let
      allFieldsC :: Cxt
allFieldsC =
        forall a b. (a -> b) -> [a] -> [b]
map (Name -> Kind
VarT Name
constraintVar Kind -> Kind -> Kind
`AppT`) (forall a b. (a -> b) -> [a] -> [b]
map 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) forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map 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
compatConP Name
name) (Exp -> Body
NormalB (Name -> Exp
ConE 'Dict)) []

    forall (f :: Type -> Type) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
      Maybe Overlap -> Cxt -> Kind -> [Dec] -> Dec
InstanceD
        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 <-
    forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap forall (t :: Type -> Type) a. Foldable t => t [a] -> [a]
concat forall a b. (a -> b) -> a -> b
$ 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 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))
        |]

  forall (f :: Type -> Type) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
    [ Dec
recordInstance
    , Dec
fieldDictInstance
    ]
    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 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

compatConP :: Name -> Pat
#if MIN_VERSION_template_haskell(2,18,0)
compatConP :: Name -> Pat
compatConP Name
constrFieldName =
    Name -> Cxt -> [Pat] -> Pat
ConP Name
constrFieldName [] []
#else
compatConP constrFieldName =
    ConP constrFieldName []
#endif