{-# 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
      case Info
info of
        TH.TyConI (TH.DataD Cxt
_ Name
_ [TyVarBndr ()]
_ Maybe Kind
_ [TH.RecC Name
_ [VarBangType]
triples] [DerivClause]
_) -> do
          [Dec]
classes <- [VarBangType] -> Q [Dec]
makeLensClasses [VarBangType]
triples
          [Dec]
instances <- Name -> [VarBangType] -> Q [Dec]
makeLensInstances Name
name [VarBangType]
triples
          forall (m :: * -> *) a. Monad m => a -> m a
return ([Dec]
classes forall a. Semigroup a => a -> a -> a
<> [Dec]
instances)
        Info
_ -> forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"reify failed"
    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 (m :: * -> *) a. Monad m => a -> m a
return []
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 (m :: * -> *) a. Monad m => a -> m a
return (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 (m :: * -> *) a. Monad m => a -> m a
return 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 (m :: * -> *) a. Monad m => a -> m a
return (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 (m :: * -> *) a. Monad m => a -> m a
return 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)))
              )
            )
            []
        ]
    ]