{-# LANGUAGE DataKinds #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE TemplateHaskell #-}

module Data.Registry.Internal.TH where

import Control.Monad.Fail (fail)
import Data.Registry.Internal.Hedgehog
import Data.Text (splitOn)
import Hedgehog
import Language.Haskell.TH
import Language.Haskell.TH.Syntax
import Protolude hiding (Type)
import Prelude (last)

-- | Create a generator for selecting between constructors of an ADT
--   One parameter is a Gen Chooser in order to be able to later on
--   switch the selection strategy
makeSelectGenerator :: Name -> [Con] -> ExpQ
makeSelectGenerator :: Name -> [Con] -> ExpQ
makeSelectGenerator Name
name [Con]
constructors = do
  Pat
chooserParam <- [p|(chooser :: Gen Chooser)|]
  [Pat]
otherParams <- (Con -> Q Pat) -> [Con] -> Q [Pat]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (Name -> Con -> Q Pat
parameterFor Name
name) [Con]
constructors
  [Exp]
untaggedGenerators <- (Con -> ExpQ) -> [Con] -> Q [Exp]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Con -> ExpQ
untagGenerator [Con]
constructors
  Exp
expression <- ExpQ -> ExpQ -> ExpQ
forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE (ExpQ -> ExpQ -> ExpQ
forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE (Name -> ExpQ
forall (m :: * -> *). Quote m => Name -> m Exp
varE (String -> Name
mkName String
"chooseOne")) (Name -> ExpQ
forall (m :: * -> *). Quote m => Name -> m Exp
varE (String -> Name
mkName String
"chooser"))) (Exp -> ExpQ
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> ExpQ) -> Exp -> ExpQ
forall a b. (a -> b) -> a -> b
$ [Exp] -> Exp
ListE [Exp]
untaggedGenerators)
  Exp -> ExpQ
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Exp -> ExpQ) -> Exp -> ExpQ
forall a b. (a -> b) -> a -> b
$ [Pat] -> Exp -> Exp
LamE (Pat
chooserParam Pat -> [Pat] -> [Pat]
forall a. a -> [a] -> [a]
: [Pat]
otherParams) Exp
expression
  where
    parameterFor :: Name -> Con -> Q Pat
    parameterFor :: Name -> Con -> Q Pat
parameterFor Name
typeName Con
constructor = do
      Name
constructorParam <- Con -> Q Name
constructorParameterName Con
constructor
      Name
constructorTag <- Con -> Q Name
tagName Con
constructor
      Q Pat -> Q Type -> Q Pat
forall (m :: * -> *). Quote m => m Pat -> m Type -> m Pat
sigP (Name -> Q Pat
forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
constructorParam) (Q Type -> Q Type -> Q Type
forall (m :: * -> *). Quote m => m Type -> m Type -> m Type
appT (Name -> Q Type
forall (m :: * -> *). Quote m => Name -> m Type
conT (String -> Name
mkName String
"Gen")) (Q Type -> Q Type -> Q Type
forall (m :: * -> *). Quote m => m Type -> m Type -> m Type
appT (Q Type -> Q Type -> Q Type
forall (m :: * -> *). Quote m => m Type -> m Type -> m Type
appT (Name -> Q Type
forall (m :: * -> *). Quote m => Name -> m Type
conT (String -> Name
mkName String
"Tag")) (Q TyLit -> Q Type
forall (m :: * -> *). Quote m => m TyLit -> m Type
litT (String -> Q TyLit
forall (m :: * -> *). Quote m => String -> m TyLit
strTyLit (Name -> String
forall a b. (Show a, StringConv String b) => a -> b
show Name
constructorTag)))) (Name -> Q Type
forall (m :: * -> *). Quote m => Name -> m Type
conT Name
typeName)))

-- Create a generator expression for a specific constructor of a data type
-- runQ [|tag @"permanent" Permanent|]
-- AppE (AppTypeE (VarE Data.Registry.Lift.tag) (LitT (StrTyLit "permanent"))) (ConE Test.Data.Registry.Generators.Permanent)
makeConstructorGenerator :: Con -> ExpQ
makeConstructorGenerator :: Con -> ExpQ
makeConstructorGenerator Con
constructor = do
  Name
constructorTag <- Con -> Q Name
tagName Con
constructor
  Name
constructorType <- Con -> Q Name
nameOf Con
constructor
  ExpQ -> ExpQ -> ExpQ
forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE (ExpQ -> Q Type -> ExpQ
forall (m :: * -> *). Quote m => m Exp -> m Type -> m Exp
appTypeE (Name -> ExpQ
forall (m :: * -> *). Quote m => Name -> m Exp
varE (String -> Name
mkName String
"tag")) (Q TyLit -> Q Type
forall (m :: * -> *). Quote m => m TyLit -> m Type
litT (String -> Q TyLit
forall (m :: * -> *). Quote m => String -> m TyLit
strTyLit (Name -> String
forall a b. (Show a, StringConv String b) => a -> b
show Name
constructorTag)))) (Name -> ExpQ
forall (m :: * -> *). Quote m => Name -> m Exp
conE Name
constructorType)

-- | Remove the tag of a given constructor: fmap unTag g :: Gen (Tag "t" SomeType) -> Gen SomeType
untagGenerator :: Con -> ExpQ
untagGenerator :: Con -> ExpQ
untagGenerator Con
constructor = do
  Name
constructorParam <- Con -> Q Name
constructorParameterName Con
constructor
  ExpQ -> ExpQ -> ExpQ
forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE (ExpQ -> ExpQ -> ExpQ
forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE (Name -> ExpQ
forall (m :: * -> *). Quote m => Name -> m Exp
varE (String -> Name
mkName String
"fmap")) (Name -> ExpQ
forall (m :: * -> *). Quote m => Name -> m Exp
varE (String -> Name
mkName String
"unTag"))) (Name -> ExpQ
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
constructorParam)

-- | Create a tag used to distinguish constructors in an ADT
tagName :: Con -> Q Name
tagName :: Con -> Q Name
tagName Con
constructor = do
  Text
name <- Con -> Q Text
constructorName Con
constructor
  Name -> Q Name
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Name -> Q Name) -> Name -> Q Name
forall a b. (a -> b) -> a -> b
$ String -> Name
mkName (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ Text -> String
forall a b. ConvertText a b => a -> b
toS Text
name

-- | Same as the tag name but lower cased
constructorParameterName :: Con -> Q Name
constructorParameterName :: Con -> Q Name
constructorParameterName Con
constructor = do
  Text
name <- Con -> Q Text
constructorName Con
constructor
  Name -> Q Name
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Name -> Q Name) -> Name -> Q Name
forall a b. (a -> b) -> a -> b
$ String -> Name
mkName (Char -> Char
toLower (Char -> Char) -> String -> String
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Text -> String
forall a b. ConvertText a b => a -> b
toS Text
name)

-- | Extract the last name of a constructor
constructorName :: Con -> Q Text
constructorName :: Con -> Q Text
constructorName Con
constructor = do
  Name
n <- Con -> Q Name
nameOf Con
constructor
  Text -> Q Text
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Text -> Q Text) -> Text -> Q Text
forall a b. (a -> b) -> a -> b
$ [Text] -> Text
forall a. [a] -> a
last ([Text] -> Text) -> (Text -> [Text]) -> Text -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Text -> [Text]
splitOn Text
"." (Text -> Text) -> Text -> Text
forall a b. (a -> b) -> a -> b
$ Name -> Text
forall a b. (Show a, StringConv String b) => a -> b
show Name
n

-- | The name of a given constructor
nameOf :: Con -> Q Name
nameOf :: Con -> Q Name
nameOf (NormalC Name
n [BangType]
_) = Name -> Q Name
forall (f :: * -> *) a. Applicative f => a -> f a
pure Name
n
nameOf (RecC Name
n [VarBangType]
_) = Name -> Q Name
forall (f :: * -> *) a. Applicative f => a -> f a
pure Name
n
nameOf Con
other = do
  Bool -> String -> Q ()
forall (m :: * -> *). Quasi m => Bool -> String -> m ()
qReport Bool
True (String
"we can only create generators for normal constructors and records, got: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Con -> String
forall a b. (Show a, StringConv String b) => a -> b
show Con
other)
  String -> Q Name
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"generators creation failed"

-- | The list of types necessary to create a given constructor
typesOf :: Con -> Q [Type]
typesOf :: Con -> Q [Type]
typesOf (NormalC Name
_ [BangType]
types) = [Type] -> Q [Type]
forall (f :: * -> *) a. Applicative f => a -> f a
pure (BangType -> Type
forall a b. (a, b) -> b
snd (BangType -> Type) -> [BangType] -> [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [BangType]
types)
typesOf (RecC Name
_ [VarBangType]
types) = [Type] -> Q [Type]
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Type] -> Q [Type]) -> [Type] -> Q [Type]
forall a b. (a -> b) -> a -> b
$ (\(Name
_, Bang
_, Type
t) -> Type
t) (VarBangType -> Type) -> [VarBangType] -> [Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [VarBangType]
types
typesOf Con
other = do
  Bool -> String -> Q ()
forall (m :: * -> *). Quasi m => Bool -> String -> m ()
qReport Bool
True (String
"we can only create generators for normal constructors and records, got: " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Con -> String
forall a b. (Show a, StringConv String b) => a -> b
show Con
other)
  String -> Q [Type]
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"generators creation failed"

-- | runQ [| fun g +: genFun e +: genFun f|]
--   InfixE (Just (AppE (VarE Data.Registry.Registry.fun) (UnboundVarE g))) (VarE +:) (Just (InfixE (Just (AppE (VarE Data.Registry.Hedgehog.genFun) (UnboundVarE e)))
--  (VarE Data.Registry.Registry.+:) (Just (AppE (VarE Data.Registry.Hedgehog.genFun) (UnboundVarE f)))))
assembleGeneratorsToRegistry :: Exp -> [Exp] -> ExpQ
assembleGeneratorsToRegistry :: Exp -> [Exp] -> ExpQ
assembleGeneratorsToRegistry Exp
_ [] =
  String -> ExpQ
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"generators creation failed"
assembleGeneratorsToRegistry Exp
selectorGenerator [Exp
g] =
  ExpQ -> ExpQ -> ExpQ
app (ExpQ -> ExpQ
genFunOf (Exp -> ExpQ
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
g)) (ExpQ -> ExpQ) -> ExpQ -> ExpQ
forall a b. (a -> b) -> a -> b
$
    ExpQ -> ExpQ -> ExpQ
app (ExpQ -> ExpQ
funOf (Exp -> ExpQ
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
selectorGenerator)) (ExpQ -> ExpQ) -> ExpQ -> ExpQ
forall a b. (a -> b) -> a -> b
$
      ExpQ -> ExpQ -> ExpQ
app (ExpQ -> ExpQ
genFunOf (Name -> ExpQ
forall (m :: * -> *). Quote m => Name -> m Exp
varE (String -> Name
mkName String
"choiceChooser"))) (Name -> ExpQ
forall (m :: * -> *). Quote m => Name -> m Exp
varE (String -> Name
mkName String
"emptyRegistry"))
--
assembleGeneratorsToRegistry Exp
selectorGenerator (Exp
g : [Exp]
gs) =
  ExpQ -> ExpQ -> ExpQ
app (ExpQ -> ExpQ
genFunOf (Exp -> ExpQ
forall (f :: * -> *) a. Applicative f => a -> f a
pure Exp
g)) (Exp -> [Exp] -> ExpQ
assembleGeneratorsToRegistry Exp
selectorGenerator [Exp]
gs)

app :: ExpQ -> ExpQ -> ExpQ
app :: ExpQ -> ExpQ -> ExpQ
app ExpQ
e1 ExpQ
e2 =
  Maybe ExpQ -> ExpQ -> Maybe ExpQ -> ExpQ
forall (m :: * -> *).
Quote m =>
Maybe (m Exp) -> m Exp -> Maybe (m Exp) -> m Exp
infixE (ExpQ -> Maybe ExpQ
forall a. a -> Maybe a
Just ExpQ
e1) (Name -> ExpQ
forall (m :: * -> *). Quote m => Name -> m Exp
varE (String -> Name
mkName String
"+:")) (ExpQ -> Maybe ExpQ
forall a. a -> Maybe a
Just ExpQ
e2)

genFunOf :: ExpQ -> ExpQ
genFunOf :: ExpQ -> ExpQ
genFunOf = ExpQ -> ExpQ -> ExpQ
forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE (Name -> ExpQ
forall (m :: * -> *). Quote m => Name -> m Exp
varE (String -> Name
mkName String
"genFun"))

funOf :: ExpQ -> ExpQ
funOf :: ExpQ -> ExpQ
funOf = ExpQ -> ExpQ -> ExpQ
forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE (Name -> ExpQ
forall (m :: * -> *). Quote m => Name -> m Exp
varE (String -> Name
mkName String
"fun"))