module Data.DeepLenses where

import Control.Lens (Lens', makeClassy)
import Language.Haskell.TH
import Language.Haskell.TH.Datatype (
  ConstructorInfo(ConstructorInfo),
  ConstructorVariant(RecordConstructor),
  DatatypeInfo(datatypeCons, datatypeName),
  reifyDatatype,
  )
import Language.Haskell.TH.Syntax (ModName(..), Name(Name), NameFlavour(NameQ, NameS, NameG), OccName(..))

class DeepLenses s s' where
  deepLens :: Lens' s s'

data Field =
  Field {
    Field -> Name
fieldName :: Name,
    Field -> Type
fieldType :: Type
  }
  deriving Int -> Field -> ShowS
[Field] -> ShowS
Field -> String
(Int -> Field -> ShowS)
-> (Field -> String) -> ([Field] -> ShowS) -> Show Field
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Field] -> ShowS
$cshowList :: [Field] -> ShowS
show :: Field -> String
$cshow :: Field -> String
showsPrec :: Int -> Field -> ShowS
$cshowsPrec :: Int -> Field -> ShowS
Show

data DT =
  DT {
    DT -> Name
dtName :: Name,
    DT -> [Field]
dtFields :: [Field]
  }
  deriving Int -> DT -> ShowS
[DT] -> ShowS
DT -> String
(Int -> DT -> ShowS)
-> (DT -> String) -> ([DT] -> ShowS) -> Show DT
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [DT] -> ShowS
$cshowList :: [DT] -> ShowS
show :: DT -> String
$cshow :: DT -> String
showsPrec :: Int -> DT -> ShowS
$cshowsPrec :: Int -> DT -> ShowS
Show

dataType :: Name -> Q DT
dataType :: Name -> Q DT
dataType Name
name = do
  DatatypeInfo
info <- Name -> Q DatatypeInfo
reifyDatatype Name
name
  return $ Name -> [Field] -> DT
DT (DatatypeInfo -> Name
datatypeName DatatypeInfo
info) ([ConstructorInfo] -> [Field]
forall l. (IsList l, Item l ~ ConstructorInfo) => l -> [Field]
fields ([ConstructorInfo] -> [Field]) -> [ConstructorInfo] -> [Field]
forall a b. (a -> b) -> a -> b
$ DatatypeInfo -> [ConstructorInfo]
datatypeCons DatatypeInfo
info)
  where
    fields :: l -> [Field]
fields [ConstructorInfo _ _ _ types _ (RecordConstructor names)] =
      (Name -> Type -> Field) -> [Name] -> [Type] -> [Field]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Name -> Type -> Field
Field [Name]
names [Type]
types
    fields l
_ =
      []

mkHoist :: TypeQ -> TypeQ -> BodyQ -> DecQ
mkHoist :: TypeQ -> TypeQ -> BodyQ -> DecQ
mkHoist TypeQ
_ TypeQ
_ BodyQ
body = do
  (VarE Name
name) <- [|deepLens|]
  Name -> [ClauseQ] -> DecQ
funD Name
name [[PatQ] -> BodyQ -> [DecQ] -> ClauseQ
clause [] BodyQ
body []]

deepLensesInstance :: TypeQ -> TypeQ -> BodyQ -> DecQ
deepLensesInstance :: TypeQ -> TypeQ -> BodyQ -> DecQ
deepLensesInstance TypeQ
top TypeQ
local' BodyQ
body =
  CxtQ -> TypeQ -> [DecQ] -> DecQ
instanceD ([TypeQ] -> CxtQ
cxt []) (TypeQ -> TypeQ -> TypeQ
appT (TypeQ -> TypeQ -> TypeQ
appT [t|DeepLenses|] TypeQ
top) TypeQ
local') [TypeQ -> TypeQ -> BodyQ -> DecQ
mkHoist TypeQ
top TypeQ
local' BodyQ
body]

idLenses :: Name -> DecQ
idLenses :: Name -> DecQ
idLenses Name
name =
  TypeQ -> TypeQ -> BodyQ -> DecQ
deepLensesInstance TypeQ
nt TypeQ
nt BodyQ
body
  where
    nt :: TypeQ
nt = Name -> TypeQ
conT Name
name
    body :: BodyQ
body = ExpQ -> BodyQ
normalB [|id|]

eligibleForDeepError :: Name -> Q Bool
eligibleForDeepError :: Name -> Q Bool
eligibleForDeepError Name
tpe = do
  (ConT Name
name) <- [t|DeepLenses|]
  Name -> [Type] -> Q Bool
isInstance Name
name [Name -> Type
ConT Name
tpe, Name -> Type
ConT Name
tpe]

modName :: NameFlavour -> Maybe ModName
modName :: NameFlavour -> Maybe ModName
modName (NameQ ModName
mod') =
  ModName -> Maybe ModName
forall a. a -> Maybe a
Just ModName
mod'
modName (NameG NameSpace
_ PkgName
_ ModName
mod') =
  ModName -> Maybe ModName
forall a. a -> Maybe a
Just ModName
mod'
modName NameFlavour
_ =
  Maybe ModName
forall a. Maybe a
Nothing

sameModule :: NameFlavour -> NameFlavour -> Bool
sameModule :: NameFlavour -> NameFlavour -> Bool
sameModule NameFlavour
f1 NameFlavour
f2 =
  case (NameFlavour -> Maybe ModName
modName NameFlavour
f1, NameFlavour -> Maybe ModName
modName NameFlavour
f2) of
    (Just ModName
a, Just ModName
b) | ModName
a ModName -> ModName -> Bool
forall a. Eq a => a -> a -> Bool
== ModName
b -> Bool
True
    (Maybe ModName, Maybe ModName)
_ -> Bool
False

lensName :: Name -> Name -> ExpQ
lensName :: Name -> Name -> ExpQ
lensName (Name OccName
_ NameFlavour
topFlavour) (Name (OccName String
n) NameFlavour
lensFlavour) =
  Name -> ExpQ
varE (OccName -> NameFlavour -> Name
Name (String -> OccName
OccName (ShowS
lensNams' String
n)) NameFlavour
flavour)
  where
    lensNams' :: ShowS
lensNams' (Char
'_' : String
t) = String
t
    lensNams' [] = []
    lensNams' String
a = String
a
    flavour :: NameFlavour
flavour
      | NameFlavour -> NameFlavour -> Bool
sameModule NameFlavour
topFlavour NameFlavour
lensFlavour = NameFlavour
NameS
      | Bool
otherwise = NameFlavour
lensFlavour

fieldLenses :: Name -> [Name] -> Field -> DecsQ
fieldLenses :: Name -> [Name] -> Field -> DecsQ
fieldLenses Name
top [Name]
intermediate (Field Name
name (ConT Name
tpe)) = do
  Dec
current <- TypeQ -> TypeQ -> BodyQ -> DecQ
deepLensesInstance (Name -> TypeQ
conT Name
top) (Name -> TypeQ
conT Name
tpe) (ExpQ -> BodyQ
normalB ExpQ
body)
  [Dec]
sub <- Name -> [Name] -> Name -> DecsQ
dataLensesIfEligible Name
top (Name
name Name -> [Name] -> [Name]
forall a. a -> [a] -> [a]
: [Name]
intermediate) Name
tpe
  return (Dec
current Dec -> [Dec] -> [Dec]
forall a. a -> [a] -> [a]
: [Dec]
sub)
  where
    compose :: Name -> ExpQ -> ExpQ
compose = ExpQ -> ExpQ -> ExpQ
appE (ExpQ -> ExpQ -> ExpQ) -> (Name -> ExpQ) -> Name -> ExpQ -> ExpQ
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ExpQ -> ExpQ -> ExpQ
appE [|(.)|] (ExpQ -> ExpQ) -> (Name -> ExpQ) -> Name -> ExpQ
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> Name -> ExpQ
lensName Name
top
    body :: ExpQ
body = (Name -> ExpQ -> ExpQ) -> ExpQ -> [Name] -> ExpQ
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Name -> ExpQ -> ExpQ
compose (Name -> Name -> ExpQ
lensName Name
top Name
name) ([Name] -> [Name]
forall a. [a] -> [a]
reverse [Name]
intermediate)
fieldLenses Name
_ [Name]
_ Field
_ =
  [Dec] -> DecsQ
forall (m :: * -> *) a. Monad m => a -> m a
return []

dataLenses :: Name -> [Name] -> Name -> DecsQ
dataLenses :: Name -> [Name] -> Name -> DecsQ
dataLenses Name
top [Name]
intermediate Name
local' = do
  (DT Name
_ [Field]
fields) <- Name -> Q DT
dataType Name
local'
  [[Dec]] -> [Dec]
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join ([[Dec]] -> [Dec]) -> Q [[Dec]] -> DecsQ
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Field -> DecsQ) -> [Field] -> Q [[Dec]]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (Name -> [Name] -> Field -> DecsQ
fieldLenses Name
top [Name]
intermediate) [Field]
fields

dataLensesIfEligible :: Name -> [Name] -> Name -> DecsQ
dataLensesIfEligible :: Name -> [Name] -> Name -> DecsQ
dataLensesIfEligible Name
top [Name]
intermediate Name
local' = do
  Bool
eligible <- Name -> Q Bool
eligibleForDeepError Name
local'
  if Bool
eligible then Name -> [Name] -> Name -> DecsQ
dataLenses Name
top [Name]
intermediate Name
local' else [Dec] -> DecsQ
forall (m :: * -> *) a. Monad m => a -> m a
return []

lensesForMainData :: Name -> DecsQ
lensesForMainData :: Name -> DecsQ
lensesForMainData Name
name = do
  Dec
idL <- Name -> DecQ
idLenses Name
name
  [Dec]
fields <- Name -> [Name] -> Name -> DecsQ
dataLenses Name
name [] Name
name
  return (Dec
idL Dec -> [Dec] -> [Dec]
forall a. a -> [a] -> [a]
: [Dec]
fields)

deepLenses :: Name -> DecsQ
deepLenses :: Name -> DecsQ
deepLenses Name
name = do
  [Dec]
lenses <- Name -> DecsQ
makeClassy Name
name
  [Dec]
err <- Name -> DecsQ
lensesForMainData Name
name
  return $ [Dec]
lenses [Dec] -> [Dec] -> [Dec]
forall a. [a] -> [a] -> [a]
++ [Dec]
err