{-# LANGUAGE
      TemplateHaskell,
      UnicodeSyntax,
      CPP
    #-}
{- |
    Exports functions for deriving instances of 'Memoizable' using
    Template Haskell.  The @TemplateHaskell@ language extension must be
    enabled to use the functions exported from this module.
-}
module Data.Function.Memoize.TH (
  deriveMemoizable, deriveMemoizableParams, deriveMemoize,
) where

#if __GLASGOW_HASKELL__ < 710
import Control.Applicative
#endif

import Control.Monad
import Language.Haskell.TH
import Data.Function.Memoize.Class

---
--- `#DEFINE`S FOR VERSION COMPATIBILITY
---

-- GHC 7.6 changed to StarT from StarK:
#if __GLASGOW_HASKELL__ >= 706
#  define COMPAT_STAR StarT
#else
#  define COMPAT_STAR StarK
#endif

--- TH 2.10 treats type classes like type constructors:
#if MIN_VERSION_template_haskell(2,10,0)
#  define COMPAT_CLASS_PRED(C)  (appT (conT (C)) . varT)
#else
#  define COMPAT_CLASS_PRED(C)  (classP (C) . (:[]) . varT)
#endif

-- TH 2.11 supports GADTs and adds a field to NewtypeD and DataD:
#if MIN_VERSION_template_haskell(2,11,0)
#  define COMPAT_TH_GADTS
#  define COMPAT_NEWTYPE_D(N, T, C)  (NewtypeD _ (N) (T) _ (C) _)
#  define COMPAT_DATA_D(N, T, C)     (DataD _ (N) (T) _ (C) _)
#else
#  undef  COMPAT_TH_GADTS
#  define COMPAT_NEWTYPE_D(N, T, C)  (NewtypeD _ (N) (T) (C) _)
#  define COMPAT_DATA_D(N, T, C)     (DataD _ (N) (T) (C) _)
#endif

-- GHC 9 adds a type parameter to the TyVarBndr type:
#if __GLASGOW_HASKELL__ >= 900
#  define COMPAT_TY_VAR_BNDR(V)      (TyVarBndr (V))
#  define COMPAT_PLAIN_TV(N)         (PlainTV (N) _)
#  define COMPAT_KINDED_TV(N, K)     (KindedTV (N) _ (K))
#else
#  define COMPAT_TY_VAR_BNDR(V)      TyVarBndr
#  define COMPAT_PLAIN_TV(N)         (PlainTV (N))
#  define COMPAT_KINDED_TV(N, K)     (KindedTV (N) (K))
#endif

-- |
-- To derive 'Memoizable' instances for the given data types.
-- In the simplest usage, to derive 'Memoizable' for an algebraic
-- datatype named @T@, write:
--
-- @
--   deriveMemoizable ''T
-- @
--
-- This assumes that all the type parameters of @T@ that are not
-- annotated with a kind other than @*@ should be listed as requiring
-- 'Memoizable' instances in the instance context.  For example, given
-- a data type declared as
--
-- @
--   data T a (b :: * -> *) c = ...
-- @
--
-- the generated instance will look like
--
-- @
--   instance ('Memoizable' a, 'Memoizable' c) =>
--            'Memoizable' (T a b c) where ...
-- @
--
-- For more precise control over the context, use
-- 'deriveMemoizableParams'.
--
-- N.B.: The @TemplateHaskell@ language extension must be enabled to use
-- this function.
deriveMemoizable ∷ Name → Q [Dec]
deriveMemoizable :: Name -> Q [Dec]
deriveMemoizable Name
n = Name -> Maybe [Int] -> Q [Dec]
deriveMemoizable' Name
n forall a. Maybe a
Nothing

-- |
-- Like 'deriveMemoizable' but takes a second argument, which is a list
-- of 'Int's to specify which type parameters of the type should be
-- mentioned in the context.  For example, given the same definition for
-- @T@ as above, we can write
--
-- @
--    deriveMemoizableParams ''T [3]
-- @
--
-- to leave the first parameter of @T@ out of the context and show
-- only the third, yielding the instance
--
-- @
--   instance 'Memoizable' c => 'Memoizable' (T a b c) where ...
-- @
--
-- N.B.: The @TemplateHaskell@ language extension must be enabled to use
-- this function.
deriveMemoizableParams ∷ Name → [Int] → Q [Dec]
deriveMemoizableParams :: Name -> [Int] -> Q [Dec]
deriveMemoizableParams Name
n [Int]
indices = Name -> Maybe [Int] -> Q [Dec]
deriveMemoizable' Name
n (forall a. a -> Maybe a
Just [Int]
indices)

-- | In cases where neither 'deriveMemoizable' nor
-- 'deriveMemoizableParams' can figure out the right context for an
-- instance declaration, one can declare the instance manually and use
-- this function to derive the method body for 'memoize'. For example,
-- suppose that a data type @T@ is defined as:
--
-- @
--   data T a b = T (a -> Bool) b
-- @
--
-- For @T a b@ to be memoizable, @a -> Bool@ must be, and based on the
-- instance for '(->)', this means that @a@ must satisfy
-- 'Bounded' and 'Enum', so 'deriveMemoizable' cannot build the right
-- context for the 'Memoizable' instance.  Instead, one can write:
--
-- @
--   instance ('Eq' a, 'Enum' a, 'Bounded' a, 'Memoizable' b) =>
--            'Memoizable' (T a b) where
--     memoize = $(deriveMemoize ''T)
-- @
deriveMemoize ∷ Name → ExpQ
deriveMemoize :: Name -> ExpQ
deriveMemoize Name
name0 = do
  (Name
_, [TyVarBndr ()]
_, [(Name, Int)]
cons) ← Name -> Q (Name, [TyVarBndr ()], [(Name, Int)])
checkName Name
name0
  [(Name, Int)] -> ExpQ
buildMethodExp [(Name, Int)]
cons

-- | The main entry point delegates to check given type name, renames type
--   parameters, and generates the instance.
deriveMemoizable' ∷ Name → Maybe [Int] → Q [Dec]
deriveMemoizable' :: Name -> Maybe [Int] -> Q [Dec]
deriveMemoizable' Name
name0 Maybe [Int]
mindices = do
  (Name
name, [TyVarBndr ()]
tvbs, [(Name, Int)]
cons) ← Name -> Q (Name, [TyVarBndr ()], [(Name, Int)])
checkName Name
name0
  let tvs :: [Name]
tvs = forall a. [a] -> [Name]
freshNames [TyVarBndr ()]
tvbs
  Dec
inst ← forall (m :: * -> *).
Quote m =>
m Cxt -> m Type -> [m Dec] -> m Dec
instanceD
           (forall a. Maybe [Int] -> [TyVarBndr a] -> [Name] -> CxtQ
buildContext Maybe [Int]
mindices [TyVarBndr ()]
tvbs [Name]
tvs)
           (Name -> [Name] -> TypeQ
buildHead Name
name [Name]
tvs)
           [[(Name, Int)] -> Q Dec
buildMethodDec [(Name, Int)]
cons]
  forall (m :: * -> *) a. Monad m => a -> m a
return [Dec
inst]

-- | Given the type name for the requested instance, checks if it
--   corresponds to a @data@ or @newtype@, and if so, returns the name,
--   a list of its parameters, and a list of constructor names with
--   their arities.
checkName ∷ Name → Q (Name, [COMPAT_TY_VAR_BNDR(())], [(Name, Int)])
checkName :: Name -> Q (Name, [TyVarBndr ()], [(Name, Int)])
checkName Name
name0 = do
  let can'tDerive :: [Char]
can'tDerive      = [Char]
"deriveMemoizable: Can’t derive a Memoizable " forall a. [a] -> [a] -> [a]
++
                         [Char]
"instance for ‘" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Name
name0 forall a. [a] -> [a] -> [a]
++ [Char]
"’ because "
      can'tDeriveNonTC :: [Char]
can'tDeriveNonTC = [Char]
can'tDerive forall a. [a] -> [a] -> [a]
++ [Char]
"it isn’t a type constructor."
      can'tDeriveGadt :: [Char]
can'tDeriveGadt  = [Char]
can'tDerive forall a. [a] -> [a] -> [a]
++ [Char]
"GADTs aren’t supported."
      --
      stdizeCon :: Con -> m (Name, Int)
stdizeCon (NormalC Name
name [BangType]
params) = forall (m :: * -> *) a. Monad m => a -> m a
return (Name
name, forall (t :: * -> *) a. Foldable t => t a -> Int
length [BangType]
params)
      stdizeCon (RecC Name
name [VarBangType]
fields)    = forall (m :: * -> *) a. Monad m => a -> m a
return (Name
name, forall (t :: * -> *) a. Foldable t => t a -> Int
length [VarBangType]
fields)
      stdizeCon (InfixC BangType
_ Name
name BangType
_)     = forall (m :: * -> *) a. Monad m => a -> m a
return (Name
name, Int
2)
      stdizeCon (ForallC [TyVarBndr Specificity]
_ Cxt
_ Con
con)     = Con -> m (Name, Int)
stdizeCon Con
con
#ifdef COMPAT_TH_GADTS
      stdizeCon (GadtC [Name]
_ [BangType]
_ Type
_)         = forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail [Char]
can'tDeriveGadt
      stdizeCon (RecGadtC [Name]
_ [VarBangType]
_ Type
_)      = forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail [Char]
can'tDeriveGadt
#endif
  --
  Info
info ← Name -> Q Info
reify Name
name0
  case Info
info of
    TyConI (COMPAT_DATA_D(name, tvbs, cons)) → do
      conInfos ← mapM stdizeCon cons
      return (name, tvbs, conInfos)
    TyConI (COMPAT_NEWTYPE_D(name, tvbs, con)) → do
      conInfo ← stdizeCon con
      return (name, tvbs, [conInfo])
    Info
_ → forall (m :: * -> *) a. MonadFail m => [Char] -> m a
fail [Char]
can'tDeriveNonTC

-- | Given a list, produces a list of nicely printable, distinct names.
--   Used so that instances print with nice parameters names, like
--
-- @
--    instance Memoizable (T a b c) where
-- @
--
-- instead of
--
-- @
--    instance Memoizable (T a[1] b[2] c32424534) where
-- @
freshNames ∷ [a] → [Name]
freshNames :: forall a. [a] -> [Name]
freshNames [a]
xs = forall a. Int -> [a] -> [a]
take (forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs) [Name]
alphabet
  where
  alphabet :: [Name]
alphabet = [ [Char] -> Name
mkName (Char
cforall a. a -> [a] -> [a]
:[Char]
s)
             | [Char]
s ← [Char]
"" forall a. a -> [a] -> [a]
: (forall a. Show a => a -> [Char]
show forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Integer
1 ∷ Integer ..])
             , Char
c ← [Char
'a' .. Char
'z'] ]

-- | Build the type class instance context, give the necessary
-- information to select which parameters to include.  If the first
-- argument is @Just ixs@, then there should be 'Memoizable' instances
-- for exactly those parameters, by index, in the context. Otherwise,
-- choose the parameters that have no explicit kind from the
-- list of binders. The third argument gives the actual type variable
-- names to use.
buildContext ∷ Maybe [Int] → [COMPAT_TY_VAR_BNDR(a)] → [Name] → CxtQ
buildContext :: forall a. Maybe [Int] -> [TyVarBndr a] -> [Name] -> CxtQ
buildContext Maybe [Int]
mindices [TyVarBndr a]
tvbs [Name]
tvs =
  forall (m :: * -> *). Quote m => [m Type] -> m Cxt
cxt (COMPAT_CLASS_PRED(''Memoizable) <$> cxttvs)
  where
  cxttvs :: [Name]
cxttvs = case Maybe [Int]
mindices of
    Just [Int]
ixs → forall a b. (a -> Bool) -> [a] -> [b] -> [b]
filterBy (forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Int]
ixs) [Int
1 ..] [Name]
tvs
    Maybe [Int]
Nothing  → forall a b. (a -> Bool) -> [a] -> [b] -> [b]
filterBy forall {flag}. TyVarBndr flag -> Bool
isStar       [TyVarBndr a]
tvbs   [Name]
tvs
  --
  isStar :: TyVarBndr flag -> Bool
isStar (TyVarBndr flag
COMPAT_PLAIN_TV(_))               = True
  isStar (COMPAT_KINDED_TV(_, COMPAT_STAR)) = True
  isStar TyVarBndr flag
_                                  = Bool
False
  --
  filterBy ∷ (a → Bool) → [a] → [b] → [b]
  filterBy :: forall a b. (a -> Bool) -> [a] -> [b] -> [b]
filterBy a -> Bool
p [a]
xs [b]
ys = forall a b. (a, b) -> b
snd forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. (a -> Bool) -> [a] -> [a]
filter (a -> Bool
p forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) (forall a b. [a] -> [b] -> [(a, b)]
zip [a]
xs [b]
ys)

-- | Build the 'Memoizable' instance head for the given type name
--   and parameter type variables.
buildHead ∷ Name → [Name] → TypeQ
buildHead :: Name -> [Name] -> TypeQ
buildHead Name
name [Name]
tvs = 
  forall (m :: * -> *). Quote m => m Type -> m Type -> m Type
appT (forall (m :: * -> *). Quote m => Name -> m Type
conT ''Memoizable) (forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl forall (m :: * -> *). Quote m => m Type -> m Type -> m Type
appT (forall (m :: * -> *). Quote m => Name -> m Type
conT Name
name) (forall (m :: * -> *). Quote m => Name -> m Type
varT forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Name]
tvs))

-- | Build the 'memoize' method. The form of 'memoize' is always
--
-- @
--      memoize f = lookup where
--        cache1 = memoize $ \x1 -> ... memoize $ \x(a1) -> f (C1 x1 ...)
--        ...
--        cacheN = memoize $ \x1 -> ... memoize $ \x(aN) -> f (CN x1 ...)
--        lookup (C1 x1 ...) = cache1 x1 ...
--        ...
--        lookup (CN xN ...) = cacheN xN ...
-- @
--
-- where @C1@ ... @CN@ are the constructors of the data type and
-- @aj@ is the arity of constructor @Cj@.
--
-- In this method, we allocate fresh names for the parameter @f@, the
-- lookup function, and the @N@ caches.  We then delegate to build
-- the definitions of @look@ and the caches.
buildMethodDec ∷ [(Name, Int)] → DecQ
buildMethodDec :: [(Name, Int)] -> Q Dec
buildMethodDec [(Name, Int)]
cons = do
  forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Dec
valD (forall (m :: * -> *). Quote m => Name -> m Pat
varP 'memoize)
    (forall (m :: * -> *). Quote m => m Exp -> m Body
normalB ([(Name, Int)] -> ExpQ
buildMethodExp [(Name, Int)]
cons))
    []

-- | Build the body of the 'memoize' method, as described in the comment
-- above 'buildMethodDec'
buildMethodExp ∷ [(Name, Int)] → ExpQ
buildMethodExp :: [(Name, Int)] -> ExpQ
buildMethodExp [(Name, Int)]
cons = do
  Name
f      ← forall (m :: * -> *). Quote m => [Char] -> m Name
newName [Char]
"fun"
  [Name]
caches ← forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\(Name, Int)
_ → forall (m :: * -> *). Quote m => [Char] -> m Name
newName [Char]
"cache") [(Name, Int)]
cons
  forall (m :: * -> *). Quote m => m Pat -> m Exp -> m Exp
lam1E (forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
f)
    (forall (m :: * -> *). Quote m => [m Dec] -> m Exp -> m Exp
letE
      (forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (Name -> (Name, Int) -> Name -> Q Dec
buildCache Name
f) [(Name, Int)]
cons [Name]
caches)
      ([(Name, Int)] -> [Name] -> ExpQ
buildLookup [(Name, Int)]
cons [Name]
caches))

-- | Build the look function by building a clause for each constructor
--   of the datatype.
buildLookup ∷ [(Name, Int)] → [Name] → ExpQ
buildLookup :: [(Name, Int)] -> [Name] -> ExpQ
buildLookup [(Name, Int)]
cons [Name]
caches = do
  Name
a ← forall (m :: * -> *). Quote m => [Char] -> m Name
newName [Char]
"arg"
  forall (m :: * -> *). Quote m => m Pat -> m Exp -> m Exp
lam1E (forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
a) forall b c a. (b -> c) -> (a -> b) -> a -> c
.
    forall (m :: * -> *). Quote m => m Exp -> [m Match] -> m Exp
caseE (forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
a) forall a b. (a -> b) -> a -> b
$
      forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (Name, Int) -> Name -> Q Match
buildLookupMatch [(Name, Int)]
cons [Name]
caches

-- | Build a lookup clause for one constructor.  We lookup a value
--   by matching that constructor and then passing its parameters to
--   the cache for that constructor.
buildLookupMatch ∷ (Name, Int) → Name → MatchQ
buildLookupMatch :: (Name, Int) -> Name -> Q Match
buildLookupMatch (Name
con, Int
arity) Name
cache = do
  [Name]
params ← forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
arity (forall (m :: * -> *). Quote m => [Char] -> m Name
newName [Char]
"param")
  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 => Name -> m Pat
varP forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Name]
params))
        (forall (m :: * -> *). Quote m => m Exp -> m Body
normalB (forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE (forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
cache) (forall (m :: * -> *). Quote m => Name -> m Exp
varE forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Name]
params)))
        []

-- | Build the definition of a cache for the given constructor.  We do
--   this by binding the cache name to a cascading sequence of
--   memoizations for each component in the constructor's arity.
buildCache ∷ Name → (Name, Int) → Name → DecQ
buildCache :: Name -> (Name, Int) -> Name -> Q Dec
buildCache Name
f (Name
con, Int
arity) Name
cache =
  forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Dec
valD (forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
cache) (forall (m :: * -> *). Quote m => m Exp -> m Body
normalB (Int -> Name -> ExpQ -> ExpQ
composeMemos Int
arity Name
f (forall (m :: * -> *). Quote m => Name -> m Exp
conE Name
con))) []

-- | Given the remaining arity to memoize, the name of the function to
--   memoize, and the accumulated parameter so far, build the
--   memoization chain.
composeMemos ∷ Int → Name → ExpQ → ExpQ
composeMemos :: Int -> Name -> ExpQ -> ExpQ
composeMemos Int
0     Name
f ExpQ
arg = [| $(varE f) $arg |]
composeMemos Int
arity Name
f ExpQ
arg = do
  [| memoize $ \b → $(composeMemos (arity - 1) f [| $arg b |]) |]