-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Singletons.Deriving.Enum
-- Copyright   :  (C) 2015 Richard Eisenberg
-- License     :  BSD-style (see LICENSE)
-- Maintainer  :  Ryan Scott
-- Stability   :  experimental
-- Portability :  non-portable
--
-- Implements deriving of Enum instances
--
----------------------------------------------------------------------------

module Data.Singletons.Deriving.Enum ( mkEnumInstance ) where

import Language.Haskell.TH.Syntax
import Language.Haskell.TH.Ppr
import Language.Haskell.TH.Desugar
import Data.Singletons.Deriving.Util
import Data.Singletons.Syntax
import Data.Singletons.Util
import Data.Singletons.Names
import Control.Monad
import Data.Maybe

-- monadic for failure only
mkEnumInstance :: DsMonad q => DerivDesc q
mkEnumInstance :: DerivDesc q
mkEnumInstance Maybe DCxt
mb_ctxt DType
ty (DataDecl Name
data_name [DTyVarBndr]
tvbs [DCon]
cons) = do
  let data_ty :: DType
data_ty = DType -> [DTyVarBndr] -> DType
foldTypeTvbs (Name -> DType
DConT Name
data_name) [DTyVarBndr]
tvbs
  Bool
non_vanilla <- DType -> [DCon] -> q Bool
forall (q :: * -> *). DsMonad q => DType -> [DCon] -> q Bool
isNonVanillaDataType DType
data_ty [DCon]
cons
  Bool -> q () -> q ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ([DCon] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [DCon]
cons Bool -> Bool -> Bool
||
        (DCon -> Bool) -> [DCon] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (\(DCon [DTyVarBndr]
_ DCxt
_ Name
_ DConFields
f DType
_) ->
              Bool
non_vanilla Bool -> Bool -> Bool
|| Bool -> Bool
not (DCxt -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (DCxt -> Bool) -> DCxt -> Bool
forall a b. (a -> b) -> a -> b
$ DConFields -> DCxt
tysOfConFields DConFields
f)) [DCon]
cons) (q () -> q ()) -> q () -> q ()
forall a b. (a -> b) -> a -> b
$
    String -> q ()
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String
"Can't derive Enum instance for " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Type -> String
forall a. Ppr a => a -> String
pprint (DType -> Type
typeToTH DType
ty) String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
".")
  Name
n <- String -> q Name
forall (m :: * -> *). Quasi m => String -> m Name
qNewName String
"n"
  let to_enum :: LetDecRHS Unannotated
to_enum = [DClause] -> LetDecRHS Unannotated
UFunction [[DPat] -> DExp -> DClause
DClause [Name -> DPat
DVarP Name
n] ([DCon] -> [Integer] -> DExp
to_enum_rhs [DCon]
cons [Integer
0..])]
      to_enum_rhs :: [DCon] -> [Integer] -> DExp
to_enum_rhs [] [Integer]
_ = Name -> DExp
DVarE Name
errorName DExp -> DExp -> DExp
`DAppE` Lit -> DExp
DLitE (String -> Lit
StringL String
"toEnum: bad argument")
      to_enum_rhs (DCon [DTyVarBndr]
_ DCxt
_ Name
name DConFields
_ DType
_ : [DCon]
rest) (Integer
num:[Integer]
nums) =
        DExp -> [DMatch] -> DExp
DCaseE (Name -> DExp
DVarE Name
equalsName DExp -> DExp -> DExp
`DAppE` Name -> DExp
DVarE Name
n DExp -> DExp -> DExp
`DAppE` Lit -> DExp
DLitE (Integer -> Lit
IntegerL Integer
num))
          [ DPat -> DExp -> DMatch
DMatch (Name -> [DPat] -> DPat
DConP Name
trueName []) (Name -> DExp
DConE Name
name)
          , DPat -> DExp -> DMatch
DMatch (Name -> [DPat] -> DPat
DConP Name
falseName []) ([DCon] -> [Integer] -> DExp
to_enum_rhs [DCon]
rest [Integer]
nums) ]
      to_enum_rhs [DCon]
_ [Integer]
_ = String -> DExp
forall a. HasCallStack => String -> a
error String
"Internal error: exhausted infinite list in to_enum_rhs"

      from_enum :: LetDecRHS Unannotated
from_enum = [DClause] -> LetDecRHS Unannotated
UFunction ((Integer -> DCon -> DClause) -> [Integer] -> [DCon] -> [DClause]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Integer
i DCon
con -> [DPat] -> DExp -> DClause
DClause [Name -> [DPat] -> DPat
DConP (DCon -> Name
extractName DCon
con) []]
                                                        (Lit -> DExp
DLitE (Integer -> Lit
IntegerL Integer
i)))
                                     [Integer
0..] [DCon]
cons)
  UInstDecl -> q UInstDecl
forall (m :: * -> *) a. Monad m => a -> m a
return (InstDecl :: forall (ann :: AnnotationFlag).
DCxt
-> Name
-> DCxt
-> OMap Name DType
-> [(Name, LetDecRHS ann)]
-> InstDecl ann
InstDecl { id_cxt :: DCxt
id_cxt     = DCxt -> Maybe DCxt -> DCxt
forall a. a -> Maybe a -> a
fromMaybe [] Maybe DCxt
mb_ctxt
                   , id_name :: Name
id_name    = Name
singletonsEnumName
                      -- need to use singletons's Enum class to get the types
                      -- to use Nat instead of Int

                   , id_arg_tys :: DCxt
id_arg_tys = [DType
ty]
                   , id_sigs :: OMap Name DType
id_sigs    = OMap Name DType
forall a. Monoid a => a
mempty
                   , id_meths :: [(Name, LetDecRHS Unannotated)]
id_meths   = [ (Name
singletonsToEnumName, LetDecRHS Unannotated
to_enum)
                                  , (Name
singletonsFromEnumName, LetDecRHS Unannotated
from_enum) ] })