-- | 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 (String -> String) -> (String -> String) -> String -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> String
dropPrefix
  }

underscore :: String -> String
underscore :: String -> String
underscore = (Char -> String) -> String -> String
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 = Int -> String -> String
forall a. Int -> [a] -> [a]
drop Int
1 (String -> String) -> (String -> String) -> String -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Bool) -> String -> String
forall a. (a -> Bool) -> [a] -> [a]
dropWhile (Char -> Char -> Bool
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
          [Dec] -> Q [Dec]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Dec]
classes [Dec] -> [Dec] -> [Dec]
forall a. [a] -> [a] -> [a]
++ [Dec]
instances)
        Info
_ -> String -> Q [Dec]
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"reify failed"
    Maybe Name
_ -> String -> Q [Dec]
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"lookupTypeName failed"

makeLensClasses :: [TH.VarStrictType] -> TH.Q [TH.Dec]
makeLensClasses :: [VarBangType] -> Q [Dec]
makeLensClasses [] = [Dec] -> Q [Dec]
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
      [Dec] -> Q [Dec]
forall (m :: * -> *) a. Monad m => a -> m a
return (Dec
klass Dec -> [Dec] -> [Dec]
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 String -> Q Dec
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"lens already exists"
    else do
      Name
a <- String -> Q Name
TH.newName String
"a"
      Name
b <- String -> Q 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)
      Dec -> Q Dec
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
  Bool -> Q Bool
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Name -> Bool
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) String -> String -> String
forall 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 Char -> String -> String
forall a. a -> [a] -> [a]
: String
cs

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

lensName :: String -> String
lensName :: String -> String
lensName String
x = if String
y String -> [String] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [String]
keywords then String
y String -> String -> String
forall 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 [VarBangType]
triples = (VarBangType -> Q Dec) -> [VarBangType] -> Q [Dec]
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) [VarBangType]
triples

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 <- String -> Q Name
TH.newName String
"f"
  Name
x <- String -> Q Name
TH.newName String
"x"
  Name
a <- String -> Q 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)
  Dec -> Q Dec
forall (m :: * -> *) a. Monad m => a -> m a
return (Dec -> Q Dec) -> Dec -> Q Dec
forall a b. (a -> b) -> a -> b
$ Maybe Overlap -> Cxt -> Kind -> [Dec] -> Dec
TH.InstanceD
    Maybe Overlap
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))))) []]]