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

Template Haskell functions for Enum types.
-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TemplateHaskell #-}

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 = [[Dec]] -> [Dec]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[Dec]] -> [Dec]) -> Q [[Dec]] -> Q [Dec]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Q [Dec]] -> Q [[Dec]]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence
  [ (Dec -> [Dec] -> [Dec]
forall a. a -> [a] -> [a]
:[]) (Dec -> [Dec]) -> Q Dec -> Q [Dec]
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' = (String -> Name) -> [String] -> [Name]
forall a b. (a -> b) -> [a] -> [b]
map String -> Name
mkName [String]
vals
    dataDec :: Q Dec
dataDec = CxtQ
-> Name
-> [TyVarBndr]
-> Maybe Kind
-> [ConQ]
-> [DerivClauseQ]
-> Q Dec
dataD ([Kind] -> CxtQ
forall (f :: * -> *) a. Applicative f => a -> f a
pure []) Name
name' [] Maybe Kind
forall a. Maybe a
Nothing ((Name -> ConQ) -> [Name] -> [ConQ]
forall a b. (a -> b) -> [a] -> [b]
map Name -> ConQ
toCon [Name]
vals') [Maybe DerivStrategy -> [PredQ] -> DerivClauseQ
derivClause Maybe DerivStrategy
forall a. Maybe a
Nothing [PredQ]
deriveClasses]
    deriveClasses :: [PredQ]
deriveClasses =
      [ [t| Eq |]
      , [t| Ord |]
      , [t| Show |]
      , [t| Enum |]
      ]
    toCon :: Name -> ConQ
toCon Name
val = Name -> [BangTypeQ] -> ConQ
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 Q [Name] -> ([Name] -> Q [Dec]) -> Q [Dec]
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 Q [Name] -> ([Name] -> Q [Dec]) -> Q [Dec]
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 = ((Dec -> Maybe Name) -> [Dec] -> [Name])
-> [Dec] -> (Dec -> Maybe Name) -> [Name]
forall a b c. (a -> b -> c) -> b -> a -> c
flip (Dec -> Maybe Name) -> [Dec] -> [Name]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe [Dec]
instances ((Dec -> Maybe Name) -> [Name]) -> (Dec -> Maybe Name) -> [Name]
forall a b. (a -> b) -> a -> b
$ \case
        InstanceD Maybe Overlap
_ [Kind]
_ (AppT Kind
_ (ConT Name
n)) [Dec]
_ -> Name -> Maybe Name
forall a. a -> Maybe a
Just Name
n
        Dec
_ -> Maybe Name
forall a. Maybe a
Nothing
  Bool -> Q () -> Q ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Name
name Name -> [Name] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Name]
instanceNames) (Q () -> Q ()) -> Q () -> Q ()
forall a b. (a -> b) -> a -> b
$ String -> Q ()
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Q ()) -> String -> Q ()
forall a b. (a -> b) -> a -> b
$ String
"Not an Enum type: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name -> String
forall a. Show a => a -> String
show Name
name

  -- extract constructor names
  Name -> Q Info
reify Name
name Q Info -> (Info -> Q [Name]) -> Q [Name]
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    TyConI (DataD [Kind]
_ Name
_ [TyVarBndr]
_ Maybe Kind
_ [Con]
cons [DerivClause]
_) -> [Con] -> (Con -> Q Name) -> Q [Name]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [Con]
cons ((Con -> Q Name) -> Q [Name]) -> (Con -> Q Name) -> Q [Name]
forall a b. (a -> b) -> a -> b
$ \case
      NormalC Name
con [] -> Name -> Q Name
forall (m :: * -> *) a. Monad m => a -> m a
return Name
con
      Con
con -> String -> Q Name
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Q Name) -> String -> Q Name
forall a b. (a -> b) -> a -> b
$ String
"Invalid constructor: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Con -> String
forall a. Show a => a -> String
show Con
con
    Info
info -> String -> Q [Name]
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Q [Name]) -> String -> Q [Name]
forall a b. (a -> b) -> a -> b
$ String
"Invalid data type: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Info -> String
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 -> PatQ
toPattern = Lit -> PatQ
litP (Lit -> PatQ) -> (Name -> Lit) -> Name -> PatQ
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Lit
stringL (String -> Lit) -> (Name -> String) -> Name -> Lit
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Char) -> String -> String
forall a b. (a -> b) -> [a] -> [b]
map Char -> Char
toLower (String -> String) -> (Name -> String) -> Name -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> String
nameBase
      toMatch :: Name -> MatchQ
toMatch Name
con = PatQ -> BodyQ -> [Q Dec] -> MatchQ
match (Name -> PatQ
toPattern Name
con) (ExpQ -> BodyQ
normalB [| pure $(conE con) |]) []

  Name
t <- String -> Q Name
newName String
"t"
  let parseEnum :: ExpQ
parseEnum = ExpQ -> [MatchQ] -> ExpQ
caseE [| Text.unpack $ Text.toLower $(varE t) |] ([MatchQ] -> ExpQ) -> [MatchQ] -> ExpQ
forall a b. (a -> b) -> a -> b
$
        (Name -> MatchQ) -> [Name] -> [MatchQ]
forall a b. (a -> b) -> [a] -> [b]
map Name -> MatchQ
toMatch [Name]
cons [MatchQ] -> [MatchQ] -> [MatchQ]
forall a. [a] -> [a] -> [a]
++ [PatQ -> BodyQ -> [Q Dec] -> MatchQ
match PatQ
wildP (ExpQ -> BodyQ
normalB (ExpQ -> BodyQ) -> ExpQ -> BodyQ
forall a b. (a -> b) -> a -> b
$ ExpQ -> ExpQ -> ExpQ
appE ExpQ
badParse (ExpQ -> ExpQ) -> ExpQ -> ExpQ
forall a b. (a -> b) -> a -> b
$ Name -> ExpQ
varE Name
t) []]

  [d|
    instance FromJSON $(conT name) where
      parseJSON (String $(varP t)) = $parseEnum
      parseJSON v = $badParse v
    |]
  where
    badParse :: ExpQ
badParse =
      let prefix :: ExpQ
prefix = Lit -> ExpQ
litE (Lit -> ExpQ) -> Lit -> ExpQ
forall a b. (a -> b) -> a -> b
$ String -> Lit
stringL (String -> Lit) -> String -> Lit
forall a b. (a -> b) -> a -> b
$ String
"Bad " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name -> String
nameBase Name
name String -> String -> String
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 -> MatchQ
encodeConstructor Name
con = PatQ -> BodyQ -> [Q Dec] -> MatchQ
match (Name -> [PatQ] -> PatQ
conP Name
con []) (ExpQ -> BodyQ
normalB [| String $ Text.pack $(lift $ nameBase con) |]) []