{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TemplateHaskell #-}

{- |
Module      :  Data.Aeson.Schema.TH.Enum
Maintainer  :  Brandon Chinn <brandon@leapyear.io>
Stability   :  experimental
Portability :  portable

Template Haskell functions for Enum types.
-}
module Data.Aeson.Schema.TH.Enum (
  genFromJSONEnum,
  genToJSONEnum,
  mkEnum,
) where

import Control.Monad (forM, unless)
import Data.Aeson (FromJSON (..), ToJSON (..), Value (..))
import Data.Char (toLower)
import Data.Maybe (mapMaybe)
import qualified Data.Text as Text
import Language.Haskell.TH
import Language.Haskell.TH.Syntax (lift)

{- | Make an enum type with the given constructors, that can be parsed from JSON.

 The 'FromJSON' instance will match to a string value matching the constructor name,
 case-insensitive.

 @
 mkEnum \"State" [\"OPEN", \"CLOSED"]

 -- generates equivalent of:
 --   data State = OPEN | CLOSED deriving (...)
 --   genFromJSONEnum ''State
 --   genToJSONEnum ''State
 @
-}
mkEnum :: String -> [String] -> Q [Dec]
mkEnum :: String -> [String] -> Q [Dec]
mkEnum String
name [String]
vals =
  forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence
      [ (forall a. a -> [a] -> [a]
: []) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Q Dec
dataDec
      , Name -> [Name] -> Q [Dec]
mkFromJSON Name
name' [Name]
vals'
      , Name -> [Name] -> Q [Dec]
mkToJSON Name
name' [Name]
vals'
      ]
  where
    name' :: Name
name' = String -> Name
mkName String
name
    vals' :: [Name]
vals' = forall a b. (a -> b) -> [a] -> [b]
map String -> Name
mkName [String]
vals
    dataDec :: Q Dec
dataDec = forall (m :: * -> *).
Quote m =>
m Cxt
-> Name
-> [TyVarBndr ()]
-> Maybe Kind
-> [m Con]
-> [m DerivClause]
-> m Dec
dataD (forall (f :: * -> *) a. Applicative f => a -> f a
pure []) Name
name' [] forall a. Maybe a
Nothing (forall a b. (a -> b) -> [a] -> [b]
map forall {m :: * -> *}. Quote m => Name -> m Con
toCon [Name]
vals') [forall (m :: * -> *).
Quote m =>
Maybe DerivStrategy -> [m Kind] -> m DerivClause
derivClause forall a. Maybe a
Nothing [Q Kind]
deriveClasses]
    deriveClasses :: [Q Kind]
deriveClasses =
      [ [t|Eq|]
      , [t|Ord|]
      , [t|Show|]
      , [t|Enum|]
      ]
    toCon :: Name -> m Con
toCon Name
val = forall (m :: * -> *). Quote m => Name -> [m BangType] -> m Con
normalC Name
val []

{- | Generate an instance of 'FromJSON' for the given data type.

 Prefer using 'mkEnum'; this function is useful for data types in which you want greater control
 over the actual data type.

 The 'FromJSON' instance will match to a string value matching the constructor name,
 case-insensitive.

 @
 data State = Open | CLOSED deriving (Show,Enum)
 genFromJSONEnum ''State

 -- outputs:
 --   Just Open
 --   Just Open
 --   Just CLOSED
 --   Just CLOSED
 main = mapM_ print
   [ decodeState \"open"
   , decodeState \"OPEN"
   , decodeState \"closed"
   , decodeState \"CLOSED"
   ]
   where
     decodeState :: String -> Maybe State
     decodeState = decode . show
 @
-}
genFromJSONEnum :: Name -> Q [Dec]
genFromJSONEnum :: Name -> Q [Dec]
genFromJSONEnum Name
name = Name -> Q [Name]
getEnumConstructors Name
name forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Name -> [Name] -> Q [Dec]
mkFromJSON Name
name

{- | Generate an instance of 'ToJSON' for the given data type.

 Prefer using 'mkEnum'; this function is useful for data types in which you want greater control
 over the actual data type.

 The 'ToJSON' instance will encode the enum as a string matching the constructor name.

 @
 data State = Open | CLOSED deriving (Show,Enum)
 genToJSONEnum ''State

 -- outputs:
 --   \"Open"
 --   \"CLOSED"
 main = mapM_ print
   [ encode Open
   , encode CLOSED
   ]
 @
-}
genToJSONEnum :: Name -> Q [Dec]
genToJSONEnum :: Name -> Q [Dec]
genToJSONEnum Name
name = Name -> Q [Name]
getEnumConstructors Name
name forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Name -> [Name] -> Q [Dec]
mkToJSON Name
name

{- Helpers -}

getEnumConstructors :: Name -> Q [Name]
getEnumConstructors :: Name -> Q [Name]
getEnumConstructors Name
name = do
  -- check if 'name' is an Enum
  ClassI Dec
_ [Dec]
instances <- Name -> Q Info
reify ''Enum
  let instanceNames :: [Name]
instanceNames = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe [Dec]
instances forall a b. (a -> b) -> a -> b
$ \case
        InstanceD Maybe Overlap
_ Cxt
_ (AppT Kind
_ (ConT Name
n)) [Dec]
_ -> forall a. a -> Maybe a
Just Name
n
        Dec
_ -> forall a. Maybe a
Nothing
  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Name
name forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Name]
instanceNames) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. MonadFail m => String -> m a
fail forall a b. (a -> b) -> a -> b
$ String
"Not an Enum type: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Name
name

  -- extract constructor names
  Name -> Q Info
reify Name
name forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    TyConI (DataD Cxt
_ Name
_ [TyVarBndr ()]
_ Maybe Kind
_ [Con]
cons [DerivClause]
_) -> forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Con]
cons forall a b. (a -> b) -> a -> b
$ \case
      NormalC Name
con [] -> forall (m :: * -> *) a. Monad m => a -> m a
return Name
con
      Con
con -> forall (m :: * -> *) a. MonadFail m => String -> m a
fail forall a b. (a -> b) -> a -> b
$ String
"Invalid constructor: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Con
con
    Info
info -> forall (m :: * -> *) a. MonadFail m => String -> m a
fail forall a b. (a -> b) -> a -> b
$ String
"Invalid data type: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Info
info

mkFromJSON :: Name -> [Name] -> Q [Dec]
mkFromJSON :: Name -> [Name] -> Q [Dec]
mkFromJSON Name
name [Name]
cons = do
  let toPattern :: Name -> Q Pat
toPattern = forall (m :: * -> *). Quote m => Lit -> m Pat
litP forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Lit
stringL forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map Char -> Char
toLower forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> String
nameBase
      toMatch :: Name -> Q Match
toMatch Name
con = forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Match
match (Name -> Q Pat
toPattern Name
con) (forall (m :: * -> *). Quote m => m Exp -> m Body
normalB [|pure $(conE con)|]) []

  Name
t <- forall (m :: * -> *). Quote m => String -> m Name
newName String
"t"
  let parseEnum :: Q Exp
parseEnum =
        forall (m :: * -> *). Quote m => m Exp -> [m Match] -> m Exp
caseE [|Text.unpack $ Text.toLower $(varE t)|] forall a b. (a -> b) -> a -> b
$
          forall a b. (a -> b) -> [a] -> [b]
map Name -> Q Match
toMatch [Name]
cons forall a. [a] -> [a] -> [a]
++ [forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Match
match forall (m :: * -> *). Quote m => m Pat
wildP (forall (m :: * -> *). Quote m => m Exp -> m Body
normalB forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE Q Exp
badParse forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
t) []]

  [d|
    instance FromJSON $(conT name) where
      parseJSON (String $(varP t)) = $parseEnum
      parseJSON v = $badParse v
    |]
  where
    badParse :: Q Exp
badParse =
      let prefix :: Q Exp
prefix = forall (m :: * -> *). Quote m => Lit -> m Exp
litE forall a b. (a -> b) -> a -> b
$ String -> Lit
stringL forall a b. (a -> b) -> a -> b
$ String
"Bad " forall a. [a] -> [a] -> [a]
++ Name -> String
nameBase Name
name forall a. [a] -> [a] -> [a]
++ String
": "
       in [|fail . ($prefix ++) . show|]

mkToJSON :: Name -> [Name] -> Q [Dec]
mkToJSON :: Name -> [Name] -> Q [Dec]
mkToJSON Name
name [Name]
cons =
  [d|
    instance ToJSON $(conT name) where
      toJSON = $(lamCaseE $ map encodeConstructor cons)
    |]
  where
    encodeConstructor :: Name -> m Match
encodeConstructor Name
con = forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Match
match (forall (m :: * -> *). Quote m => Name -> [m Pat] -> m Pat
conP Name
con []) (forall (m :: * -> *). Quote m => m Exp -> m Body
normalB [|String $ Text.pack $(lift $ nameBase con)|]) []