{-# LANGUAGE LambdaCase, RecordWildCards, ViewPatterns #-}

module Data.MakeEnum(
    makeEnum,
    makeEnumWith,
  ) where

import Control.Monad
import Data.Functor.Identity
import Data.Maybe
import Data.Monoid
import Language.Haskell.TH
import Language.Haskell.TH.Syntax
import Lens.Micro hiding(filtered)

import Data.MakeEnum.Options

makeEnum :: Name -> [Name] -> Q [Dec]
makeEnum :: Name -> [Name] -> Q [Dec]
makeEnum Name
tyName [Name]
omit = Name -> [Name] -> Options -> Q [Dec]
makeEnumWith Name
tyName [Name]
omit Options
defaultOptions

makeEnumWith :: Name -> [Name] -> Options -> Q [Dec]
makeEnumWith :: Name -> [Name] -> Options -> Q [Dec]
makeEnumWith Name
tyName [Name]
omit Options
options = Name -> Q Info
reify Name
tyName forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
  TyConI (Dec -> Maybe DataDef
unwrapDec -> Just DataDef
dec) -> do
    let deducedOpts :: DeducedOptions
deducedOpts = DataDef -> Options -> DeducedOptions
deduceOptions DataDef
dec Options
options
    let (Dec
dec', [Con]
origCons, Name
name) = DeducedOptions -> [Maybe Name] -> DataDef -> (Dec, [Con], Name)
buildReducedEnum DeducedOptions
deducedOpts [Maybe Name]
omit' DataDef
dec
    (Dec
fromSig, Dec
fromFun) <- DeducedOptions -> Name -> [Con] -> Q (Dec, Dec)
buildFromFun DeducedOptions
deducedOpts Name
name [Con]
origCons
    (Dec
toSig, Dec
toFun) <- DeducedOptions -> Name -> [Con] -> Q (Dec, Dec)
buildToFun DeducedOptions
deducedOpts Name
name [Con]
origCons
    forall (f :: * -> *) a. Applicative f => a -> f a
pure [Dec
dec', Dec
fromSig, Dec
fromFun, Dec
toSig, Dec
toFun]
  Info
_ -> forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"unsupported type"
  where omit' :: [Maybe Name]
omit' = forall a. a -> Maybe a
Just forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Name]
omit

data DataDef = DataDef Cxt Name [TyVarBndr ()] (Maybe Kind) [Con] [DerivClause]
             deriving (DataDef -> DataDef -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: DataDef -> DataDef -> Bool
$c/= :: DataDef -> DataDef -> Bool
== :: DataDef -> DataDef -> Bool
$c== :: DataDef -> DataDef -> Bool
Eq, Eq DataDef
DataDef -> DataDef -> Bool
DataDef -> DataDef -> Ordering
DataDef -> DataDef -> DataDef
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: DataDef -> DataDef -> DataDef
$cmin :: DataDef -> DataDef -> DataDef
max :: DataDef -> DataDef -> DataDef
$cmax :: DataDef -> DataDef -> DataDef
>= :: DataDef -> DataDef -> Bool
$c>= :: DataDef -> DataDef -> Bool
> :: DataDef -> DataDef -> Bool
$c> :: DataDef -> DataDef -> Bool
<= :: DataDef -> DataDef -> Bool
$c<= :: DataDef -> DataDef -> Bool
< :: DataDef -> DataDef -> Bool
$c< :: DataDef -> DataDef -> Bool
compare :: DataDef -> DataDef -> Ordering
$ccompare :: DataDef -> DataDef -> Ordering
Ord, Int -> DataDef -> ShowS
[DataDef] -> ShowS
DataDef -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [DataDef] -> ShowS
$cshowList :: [DataDef] -> ShowS
show :: DataDef -> String
$cshow :: DataDef -> String
showsPrec :: Int -> DataDef -> ShowS
$cshowsPrec :: Int -> DataDef -> ShowS
Show)

unwrapDec :: Dec -> Maybe DataDef
unwrapDec :: Dec -> Maybe DataDef
unwrapDec (DataD Cxt
cx Name
name [TyVarBndr ()]
bndrs Maybe Kind
kind [Con]
cons [DerivClause]
derivs) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Cxt
-> Name
-> [TyVarBndr ()]
-> Maybe Kind
-> [Con]
-> [DerivClause]
-> DataDef
DataDef Cxt
cx Name
name [TyVarBndr ()]
bndrs Maybe Kind
kind [Con]
cons [DerivClause]
derivs
unwrapDec Dec
_ = forall a. Maybe a
Nothing

type DeducedOptions = OptionsT Identity

deduceOptions :: DataDef -> Options -> DeducedOptions
deduceOptions :: DataDef -> Options -> DeducedOptions
deduceOptions (DataDef Cxt
_ Name
name [TyVarBndr ()]
_ Maybe Kind
_ [Con]
_ [DerivClause]
_) Options { [Name]
Maybe String
ShowS
deriveClasses :: forall (f :: * -> *). OptionsT f -> [Name]
ctorNameModifier :: forall (f :: * -> *). OptionsT f -> ShowS
toFunctionName :: forall (f :: * -> *). OptionsT f -> f String
fromFunctionName :: forall (f :: * -> *). OptionsT f -> f String
newEnumName :: forall (f :: * -> *). OptionsT f -> f String
deriveClasses :: [Name]
ctorNameModifier :: ShowS
toFunctionName :: Maybe String
fromFunctionName :: Maybe String
newEnumName :: Maybe String
.. } =
  Options
    { newEnumName :: Identity String
newEnumName = forall a. a -> Identity a
Identity forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a -> a
fromMaybe (Name -> String
nameBase Name
name) Maybe String
newEnumName
    , fromFunctionName :: Identity String
fromFunctionName = forall a. a -> Identity a
Identity forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a -> a
fromMaybe (String
"from" forall a. Semigroup a => a -> a -> a
<> Name -> String
nameBase Name
name) Maybe String
fromFunctionName
    , toFunctionName :: Identity String
toFunctionName = forall a. a -> Identity a
Identity forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a -> a
fromMaybe (String
"to" forall a. Semigroup a => a -> a -> a
<> Name -> String
nameBase Name
name) Maybe String
toFunctionName
    , [Name]
ShowS
deriveClasses :: [Name]
ctorNameModifier :: ShowS
deriveClasses :: [Name]
ctorNameModifier :: ShowS
..
    }

buildReducedEnum :: DeducedOptions -> [Maybe Name] -> DataDef -> (Dec, [Con], Name)
buildReducedEnum :: DeducedOptions -> [Maybe Name] -> DataDef -> (Dec, [Con], Name)
buildReducedEnum Options { [Name]
Identity String
ShowS
deriveClasses :: [Name]
ctorNameModifier :: ShowS
toFunctionName :: Identity String
fromFunctionName :: Identity String
newEnumName :: Identity String
deriveClasses :: forall (f :: * -> *). OptionsT f -> [Name]
ctorNameModifier :: forall (f :: * -> *). OptionsT f -> ShowS
toFunctionName :: forall (f :: * -> *). OptionsT f -> f String
fromFunctionName :: forall (f :: * -> *). OptionsT f -> f String
newEnumName :: forall (f :: * -> *). OptionsT f -> f String
.. } [Maybe Name]
omit (DataDef Cxt
cx Name
name [TyVarBndr ()]
bndrs Maybe Kind
kind [Con]
cons [DerivClause]
_) = (Cxt
-> Name
-> [TyVarBndr ()]
-> Maybe Kind
-> [Con]
-> [DerivClause]
-> Dec
DataD Cxt
cx Name
name' [TyVarBndr ()]
bndrs Maybe Kind
kind [Con]
cons' [DerivClause]
derivs, [Con]
filtered, Name
name)
  where filtered :: [Con]
filtered = forall a. (a -> Bool) -> [a] -> [a]
filter ((forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [Maybe Name]
omit) forall b c a. (b -> c) -> (a -> b) -> a -> c
. (forall s a. s -> Getting (First a) s a -> Maybe a
^? Traversal' Con Name
nameT)) [Con]
cons
        cons' :: [Con]
cons' = Traversal' Con Name
nameT forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
`over` (String -> Name
mkName forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShowS
ctorNameModifier forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> String
nameBase) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Con]
filtered
        name' :: Name
name' = String -> Name
mkName forall a b. (a -> b) -> a -> b
$ forall a. Identity a -> a
runIdentity Identity String
newEnumName
        derivs :: [DerivClause]
derivs = [Maybe DerivStrategy -> Cxt -> DerivClause
DerivClause forall a. Maybe a
Nothing forall a b. (a -> b) -> a -> b
$ Name -> Kind
ConT forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Name]
deriveClasses]

buildFromFun :: DeducedOptions -> Name -> [Con] -> Q (Dec, Dec)
buildFromFun :: DeducedOptions -> Name -> [Con] -> Q (Dec, Dec)
buildFromFun Options { [Name]
Identity String
ShowS
deriveClasses :: [Name]
ctorNameModifier :: ShowS
toFunctionName :: Identity String
fromFunctionName :: Identity String
newEnumName :: Identity String
deriveClasses :: forall (f :: * -> *). OptionsT f -> [Name]
ctorNameModifier :: forall (f :: * -> *). OptionsT f -> ShowS
toFunctionName :: forall (f :: * -> *). OptionsT f -> f String
fromFunctionName :: forall (f :: * -> *). OptionsT f -> f String
newEnumName :: forall (f :: * -> *). OptionsT f -> f String
.. } Name
name [Con]
cons = do
  Module PkgName
_ (ModName String
thisModName) <- Q Module
thisModule

  let funName :: Name
funName = String -> Name
mkName forall a b. (a -> b) -> a -> b
$ forall a. Identity a -> a
runIdentity Identity String
fromFunctionName
  let funSig :: Dec
funSig = Name -> Kind -> Dec
SigD Name
funName forall a b. (a -> b) -> a -> b
$ Kind
ArrowT Kind -> Kind -> Kind
`AppT` Name -> Kind
ConT Name
name Kind -> Kind -> Kind
`AppT` (Name -> Kind
ConT (String -> Name
mkName String
"Maybe") Kind -> Kind -> Kind
`AppT` Name -> Kind
ConT (String -> Name
mkName forall a b. (a -> b) -> a -> b
$ forall a. Identity a -> a
runIdentity Identity String
newEnumName))

  [Clause]
clauses <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall {m :: * -> *}.
(Quote m, MonadFail m) =>
String -> Con -> m Clause
mkClause String
thisModName) [Con]
cons
  let fallback :: Clause
fallback = [Pat] -> Body -> [Dec] -> Clause
Clause [Pat
WildP] (Exp -> Body
NormalB forall a b. (a -> b) -> a -> b
$ Name -> Exp
ConE forall a b. (a -> b) -> a -> b
$ String -> Name
mkName String
"Nothing") []
  let funDef :: Dec
funDef = Name -> [Clause] -> Dec
FunD Name
funName forall a b. (a -> b) -> a -> b
$ [Clause]
clauses forall a. [a] -> [a] -> [a]
++ [Clause
fallback]

  forall (f :: * -> *) a. Applicative f => a -> f a
pure (Dec
funSig, Dec
funDef)

  where
    mkClause :: String -> Con -> m Clause
mkClause String
thisModName (NormalC Name
n [BangType]
ts) = do
      let thisName :: Name
thisName = String -> Name
mkName forall a b. (a -> b) -> a -> b
$ String
thisModName forall a. Semigroup a => a -> a -> a
<> String
"." forall a. Semigroup a => a -> a -> a
<> ShowS
ctorNameModifier (Name -> String
nameBase Name
n)
      [Name]
binders <- forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM (forall (t :: * -> *) a. Foldable t => t a -> Int
length [BangType]
ts) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). Quote m => String -> m Name
newName String
"p"
      let body :: Body
body = Exp -> Body
NormalB forall a b. (a -> b) -> a -> b
$ Name -> Exp
ConE (String -> Name
mkName String
"Just") Exp -> Exp -> Exp
`AppE` (Name -> Exp
ConE Name
thisName Exp -> [Name] -> Exp
`foldBinders` [Name]
binders)
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [Pat] -> Body -> [Dec] -> Clause
Clause [Name -> Cxt -> [Pat] -> Pat
ConP Name
n [] forall a b. (a -> b) -> a -> b
$ Name -> Pat
VarP forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Name]
binders] Body
body []
    mkClause String
_ Con
p = forall (m :: * -> *) a. MonadFail m => String -> m a
fail forall a b. (a -> b) -> a -> b
$ String
"this type of constructor is not supported yet:\n" forall a. Semigroup a => a -> a -> a
<> forall a. Ppr a => a -> String
pprint Con
p

buildToFun :: DeducedOptions -> Name -> [Con] -> Q (Dec, Dec)
buildToFun :: DeducedOptions -> Name -> [Con] -> Q (Dec, Dec)
buildToFun Options { [Name]
Identity String
ShowS
deriveClasses :: [Name]
ctorNameModifier :: ShowS
toFunctionName :: Identity String
fromFunctionName :: Identity String
newEnumName :: Identity String
deriveClasses :: forall (f :: * -> *). OptionsT f -> [Name]
ctorNameModifier :: forall (f :: * -> *). OptionsT f -> ShowS
toFunctionName :: forall (f :: * -> *). OptionsT f -> f String
fromFunctionName :: forall (f :: * -> *). OptionsT f -> f String
newEnumName :: forall (f :: * -> *). OptionsT f -> f String
.. } Name
name [Con]
cons = do
  Module PkgName
_ (ModName String
thisModName) <- Q Module
thisModule

  let funName :: Name
funName = String -> Name
mkName forall a b. (a -> b) -> a -> b
$ forall a. Identity a -> a
runIdentity Identity String
toFunctionName
  let funSig :: Dec
funSig = Name -> Kind -> Dec
SigD Name
funName forall a b. (a -> b) -> a -> b
$ Kind
ArrowT Kind -> Kind -> Kind
`AppT` Name -> Kind
ConT (String -> Name
mkName forall a b. (a -> b) -> a -> b
$ forall a. Identity a -> a
runIdentity Identity String
newEnumName) Kind -> Kind -> Kind
`AppT` Name -> Kind
ConT Name
name

  [Clause]
clauses <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (forall {m :: * -> *}.
(Quote m, MonadFail m) =>
String -> Con -> m Clause
mkClause String
thisModName) [Con]
cons
  let funDef :: Dec
funDef = Name -> [Clause] -> Dec
FunD Name
funName [Clause]
clauses

  forall (f :: * -> *) a. Applicative f => a -> f a
pure (Dec
funSig, Dec
funDef)

  where
    mkClause :: String -> Con -> m Clause
mkClause String
thisModName (NormalC Name
n [BangType]
ts) = do
      let thisName :: Name
thisName = String -> Name
mkName forall a b. (a -> b) -> a -> b
$ String
thisModName forall a. Semigroup a => a -> a -> a
<> String
"." forall a. Semigroup a => a -> a -> a
<> ShowS
ctorNameModifier (Name -> String
nameBase Name
n)
      [Name]
binders <- forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM (forall (t :: * -> *) a. Foldable t => t a -> Int
length [BangType]
ts) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). Quote m => String -> m Name
newName String
"p"
      let body :: Body
body = Exp -> Body
NormalB forall a b. (a -> b) -> a -> b
$ Name -> Exp
ConE Name
n Exp -> [Name] -> Exp
`foldBinders` [Name]
binders
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [Pat] -> Body -> [Dec] -> Clause
Clause [Name -> Cxt -> [Pat] -> Pat
ConP Name
thisName [] forall a b. (a -> b) -> a -> b
$ Name -> Pat
VarP forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Name]
binders] Body
body []
    mkClause String
_ Con
p = forall (m :: * -> *) a. MonadFail m => String -> m a
fail forall a b. (a -> b) -> a -> b
$ String
"this type of constructor is not supported yet:\n" forall a. Semigroup a => a -> a -> a
<> forall a. Ppr a => a -> String
pprint Con
p

foldBinders :: Exp -> [Name] -> Exp
foldBinders :: Exp -> [Name] -> Exp
foldBinders Exp
name = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Exp -> Exp -> Exp
AppE Exp
name forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map Name -> Exp
VarE

nameT :: Traversal' Con Name
nameT :: Traversal' Con Name
nameT Name -> f Name
f (NormalC Name
n [BangType]
bts) = (Name -> [BangType] -> Con
`NormalC` [BangType]
bts) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> f Name
f Name
n
nameT Name -> f Name
f (RecC Name
n [VarBangType]
vbts) = (Name -> [VarBangType] -> Con
`RecC` [VarBangType]
vbts) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> f Name
f Name
n
nameT Name -> f Name
f (InfixC BangType
bt1 Name
n BangType
bt2) = (\Name
n' -> BangType -> Name -> BangType -> Con
InfixC BangType
bt1 Name
n' BangType
bt2) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Name -> f Name
f Name
n
nameT Name -> f Name
f (ForallC [TyVarBndr Specificity]
tvbs Cxt
cx Con
n) = [TyVarBndr Specificity] -> Cxt -> Con -> Con
ForallC [TyVarBndr Specificity]
tvbs Cxt
cx forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Traversal' Con Name
nameT Name -> f Name
f Con
n
nameT Name -> f Name
_ c :: Con
c@GadtC {} = forall (f :: * -> *) a. Applicative f => a -> f a
pure Con
c
nameT Name -> f Name
_ c :: Con
c@RecGadtC {} = forall (f :: * -> *) a. Applicative f => a -> f a
pure Con
c