{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE TemplateHaskell #-}

module DeriveHasField (
  module GHC.Records,
  deriveHasFieldWith,
)
where

import Control.Monad
import Data.Char (toLower)
import Data.Foldable as Foldable
import Data.Traversable (for)
import GHC.Records
import Language.Haskell.TH
import Language.Haskell.TH.Datatype

deriveHasFieldWith :: (String -> String) -> Name -> DecsQ
deriveHasFieldWith :: (String -> String) -> Name -> DecsQ
deriveHasFieldWith String -> String
fieldModifier = (String -> String) -> DatatypeInfo -> DecsQ
makeDeriveHasField String -> String
fieldModifier forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< Name -> Q DatatypeInfo
reifyDatatype

makeDeriveHasField :: (String -> String) -> DatatypeInfo -> DecsQ
makeDeriveHasField :: (String -> String) -> DatatypeInfo -> DecsQ
makeDeriveHasField String -> String
fieldModifier DatatypeInfo
datatypeInfo = do
  -- We do not support sum of product types
  ConstructorInfo
constructorInfo <- case DatatypeInfo
datatypeInfo.datatypeCons of
    [ConstructorInfo
info] -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ConstructorInfo
info
    [ConstructorInfo]
_ -> forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"deriveHasField: only supports product types with a single data constructor"

  -- We only support data and newtype declarations
  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (DatatypeInfo
datatypeInfo.datatypeVariant forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`Foldable.notElem` [DatatypeVariant
Datatype, DatatypeVariant
Newtype]) forall a b. (a -> b) -> a -> b
$
    forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"deriveHasField: only supports data and newtype"

  -- We only support data types with field names and concrete types
  let isConcreteType :: Type -> Bool
isConcreteType = \case
        ConT Name
_ -> Bool
True
        AppT Type
_ Type
_ -> Bool
True
        Type
_ -> Bool
False
  [Name]
recordConstructorNames <- case ConstructorInfo
constructorInfo.constructorVariant of
    RecordConstructor [Name]
names -> forall (f :: * -> *) a. Applicative f => a -> f a
pure [Name]
names
    ConstructorVariant
_ -> forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"deriveHasField: only supports constructors with field names"
  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
Foldable.all Type -> Bool
isConcreteType ConstructorInfo
constructorInfo.constructorFields) forall a b. (a -> b) -> a -> b
$
    forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"deriveHasField: only supports concrete field types"

  -- Build the instances
  let constructorNamesAndTypes :: [(Name, Type)]
      constructorNamesAndTypes :: [(Name, Type)]
constructorNamesAndTypes = forall a b. [a] -> [b] -> [(a, b)]
zip [Name]
recordConstructorNames ConstructorInfo
constructorInfo.constructorFields
  [Decs]
decs <- forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for [(Name, Type)]
constructorNamesAndTypes forall a b. (a -> b) -> a -> b
$ \(Name
name, Type
ty) ->
    let currentFieldName :: String
currentFieldName = Name -> String
nameBase Name
name
        wantedFieldName :: String
wantedFieldName = String -> String
lowerFirst forall a b. (a -> b) -> a -> b
$ String -> String
fieldModifier String
currentFieldName
        litTCurrentField :: Q Type
litTCurrentField = forall (m :: * -> *). Quote m => m TyLit -> m Type
litT forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). Quote m => String -> m TyLit
strTyLit String
currentFieldName
        litTFieldWanted :: Q Type
litTFieldWanted = forall (m :: * -> *). Quote m => m TyLit -> m Type
litT forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). Quote m => String -> m TyLit
strTyLit String
wantedFieldName
        parentTypeConstructor :: Q Type
parentTypeConstructor = forall (m :: * -> *). Quote m => Name -> m Type
conT DatatypeInfo
datatypeInfo.datatypeName
     in if String
currentFieldName forall a. Eq a => a -> a -> Bool
== String
wantedFieldName
          then forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"deriveHasField: after applying fieldModifier, field didn't change"
          else
            [d|
              instance HasField $litTFieldWanted $parentTypeConstructor $(pure ty) where
                getField = $(appTypeE (varE $ mkName "getField") litTCurrentField)
              |]
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) a. Foldable t => t [a] -> [a]
Foldable.concat [Decs]
decs

lowerFirst :: String -> String
lowerFirst :: String -> String
lowerFirst = \case
  [] -> []
  (Char
x : String
xs) -> Char -> Char
toLower Char
x forall a. a -> [a] -> [a]
: String
xs