{-# LANGUAGE CPP #-}

-- | Helper functions for template Haskell, to avoid stage restrictions.
module Strive.Internal.TH
  ( options,
    makeLenses,
  )
where

import Data.Aeson.TH (Options, defaultOptions, fieldLabelModifier)
import Data.Char (isUpper, toLower, toUpper)
import Data.Maybe (isJust)
import qualified Language.Haskell.TH as TH
import qualified Language.Haskell.TH.Syntax as TH

-- | Default FromJSON options.
options :: Options
options :: Options
options = Options
defaultOptions {fieldLabelModifier :: String -> String
fieldLabelModifier = String -> String
underscore forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> String
dropPrefix}

underscore :: String -> String
underscore :: String -> String
underscore = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Char -> String
go
  where
    go :: Char -> String
go Char
c = if Char -> Bool
isUpper Char
c then [Char
'_', Char -> Char
toLower Char
c] else [Char
c]

dropPrefix :: String -> String
dropPrefix :: String -> String
dropPrefix = forall a. Int -> [a] -> [a]
drop Int
1 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> Bool) -> [a] -> [a]
dropWhile (forall a. Eq a => a -> a -> Bool
/= Char
'_')

-- | Generate lens classes and instances for a type.
makeLenses :: String -> TH.Q [TH.Dec]
makeLenses :: String -> Q [Dec]
makeLenses String
string = do
  Maybe Name
maybeName <- String -> Q (Maybe Name)
TH.lookupTypeName String
string
  case Maybe Name
maybeName of
    Just Name
name -> do
      Info
info <- Name -> Q Info
TH.reify Name
name
      [VarBangType]
triples <- case Info
info of
        TH.TyConI (TH.DataD Cxt
_ Name
_ [TyVarBndr ()]
_ Maybe Kind
_ [TH.RecC Name
_ [VarBangType]
x] [DerivClause]
_) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure [VarBangType]
x
        TH.TyConI (TH.NewtypeD Cxt
_ Name
_ [TyVarBndr ()]
_ Maybe Kind
_ (TH.RecC Name
_ [VarBangType]
x) [DerivClause]
_) -> forall (f :: * -> *) a. Applicative f => a -> f a
pure [VarBangType]
x
        Info
_ -> forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"reify failed"
      [Dec]
classes <- [VarBangType] -> Q [Dec]
makeLensClasses [VarBangType]
triples
      [Dec]
instances <- Name -> [VarBangType] -> Q [Dec]
makeLensInstances Name
name [VarBangType]
triples
      forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Dec]
classes forall a. Semigroup a => a -> a -> a
<> [Dec]
instances)
    Maybe Name
_ -> forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"lookupTypeName failed"

makeLensClasses :: [TH.VarStrictType] -> TH.Q [TH.Dec]
makeLensClasses :: [VarBangType] -> Q [Dec]
makeLensClasses [] = forall (f :: * -> *) a. Applicative f => a -> f a
pure []
makeLensClasses (VarBangType
triple : [VarBangType]
triples) = do
  Bool
exists <- VarBangType -> Q Bool
lensExists VarBangType
triple
  if Bool
exists
    then [VarBangType] -> Q [Dec]
makeLensClasses [VarBangType]
triples
    else do
      Dec
klass <- VarBangType -> Q Dec
makeLensClass VarBangType
triple
      [Dec]
classes <- [VarBangType] -> Q [Dec]
makeLensClasses [VarBangType]
triples
      forall (f :: * -> *) a. Applicative f => a -> f a
pure (Dec
klass forall a. a -> [a] -> [a]
: [Dec]
classes)

makeLensClass :: TH.VarStrictType -> TH.Q TH.Dec
makeLensClass :: VarBangType -> Q Dec
makeLensClass VarBangType
triple = do
  Bool
exists <- VarBangType -> Q Bool
lensExists VarBangType
triple
  if Bool
exists
    then forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"lens already exists"
    else do
      Name
a <- forall (m :: * -> *). Quote m => String -> m Name
TH.newName String
"a"
      Name
b <- forall (m :: * -> *). Quote m => String -> m Name
TH.newName String
"b"
      let klass :: Dec
klass = Cxt -> Name -> [TyVarBndr ()] -> [FunDep] -> [Dec] -> Dec
TH.ClassD [] Name
name [TyVarBndr ()]
types [FunDep]
dependencies [Dec]
declarations
          name :: Name
name = String -> Name
TH.mkName (VarBangType -> String
getLensName VarBangType
triple)
          types :: [TyVarBndr ()]
types = [Name -> TyVarBndr ()
TH.plainTV Name
a, Name -> TyVarBndr ()
TH.plainTV Name
b]
          dependencies :: [FunDep]
dependencies = [[Name] -> [Name] -> FunDep
TH.FunDep [Name
a] [Name
b]]
          declarations :: [Dec]
declarations = [Name -> Kind -> Dec
TH.SigD Name
field Kind
typ]
          field :: Name
field = String -> Name
TH.mkName (VarBangType -> String
getFieldName VarBangType
triple)
          typ :: Kind
typ =
            Kind -> Kind -> Kind
TH.AppT
              (Kind -> Kind -> Kind
TH.AppT (Name -> Kind
TH.ConT (String -> Name
TH.mkName String
"Lens")) (Name -> Kind
TH.VarT Name
a))
              (Name -> Kind
TH.VarT Name
b)
      forall (f :: * -> *) a. Applicative f => a -> f a
pure Dec
klass

lensExists :: TH.VarStrictType -> TH.Q Bool
lensExists :: VarBangType -> Q Bool
lensExists VarBangType
triple = do
  let name :: String
name = VarBangType -> String
getLensName VarBangType
triple
  Maybe Name
maybeName <- String -> Q (Maybe Name)
TH.lookupTypeName String
name
  forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. Maybe a -> Bool
isJust Maybe Name
maybeName)

getLensName :: TH.VarStrictType -> String
getLensName :: VarBangType -> String
getLensName VarBangType
triple = String -> String
capitalize (VarBangType -> String
getFieldName VarBangType
triple) forall a. Semigroup a => a -> a -> a
<> String
"Lens"

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

getFieldName :: TH.VarStrictType -> String
getFieldName :: VarBangType -> String
getFieldName (Name
var, Bang
_, Kind
_) = (String -> String
lensName forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => a -> String
show) Name
var

lensName :: String -> String
lensName :: String -> String
lensName String
x = if String
y forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [String]
keywords then String
y forall a. Semigroup a => a -> a -> a
<> String
"_" else String
y
  where
    y :: String
y = String -> String
dropPrefix String
x
    keywords :: [String]
keywords = [String
"data", String
"type"]

makeLensInstances :: TH.Name -> [TH.VarStrictType] -> TH.Q [TH.Dec]
makeLensInstances :: Name -> [VarBangType] -> Q [Dec]
makeLensInstances Name
name = forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Name -> VarBangType -> Q Dec
makeLensInstance Name
name)

makeLensInstance :: TH.Name -> TH.VarStrictType -> TH.Q TH.Dec
makeLensInstance :: Name -> VarBangType -> Q Dec
makeLensInstance Name
name triple :: VarBangType
triple@(Name
var, Bang
_, Kind
typ) = do
  Name
f <- forall (m :: * -> *). Quote m => String -> m Name
TH.newName String
"f"
  Name
x <- forall (m :: * -> *). Quote m => String -> m Name
TH.newName String
"x"
  Name
a <- forall (m :: * -> *). Quote m => String -> m Name
TH.newName String
"a"
  Just Name
fmap' <- String -> Q (Maybe Name)
TH.lookupValueName String
"fmap"
  let field :: Name
field = String -> Name
TH.mkName (VarBangType -> String
getFieldName VarBangType
triple)
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$
    Maybe Overlap -> Cxt -> Kind -> [Dec] -> Dec
TH.InstanceD
      forall a. Maybe a
Nothing
      []
      ( Kind -> Kind -> Kind
TH.AppT
          (Kind -> Kind -> Kind
TH.AppT (Name -> Kind
TH.ConT (String -> Name
TH.mkName (VarBangType -> String
getLensName VarBangType
triple))) (Name -> Kind
TH.ConT Name
name))
          Kind
typ
      )
      [ Name -> [Clause] -> Dec
TH.FunD
          Name
field
          [ [Pat] -> Body -> [Dec] -> Clause
TH.Clause
              [Name -> Pat
TH.VarP Name
f, Name -> Pat
TH.VarP Name
x]
              ( Exp -> Body
TH.NormalB
                  ( Exp -> Exp -> Exp
TH.AppE
                      ( Exp -> Exp -> Exp
TH.AppE
                          (Name -> Exp
TH.VarE Name
fmap')
                          ( [Pat] -> Exp -> Exp
TH.LamE
                              [Name -> Pat
TH.VarP Name
a]
                              (Exp -> [FieldExp] -> Exp
TH.RecUpdE (Name -> Exp
TH.VarE Name
x) [(Name
var, Name -> Exp
TH.VarE Name
a)])
                          )
                      )
                      (Exp -> Exp -> Exp
TH.AppE (Name -> Exp
TH.VarE Name
f) (Exp -> Exp -> Exp
TH.AppE (Name -> Exp
TH.VarE Name
var) (Name -> Exp
TH.VarE Name
x)))
                  )
              )
              []
          ]
      ]