{-# 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

-- |
-- 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 n :: Name
n = Name -> Maybe [Int] -> Q [Dec]
deriveMemoizable' Name
n Maybe [Int]
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 n :: Name
n indices :: [Int]
indices = Name -> Maybe [Int] -> Q [Dec]
deriveMemoizable' Name
n ([Int] -> Maybe [Int]
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 ('Enum' a, 'Bounded' a, 'Memoizable' b) =>
--            'Memoizable' (T a b) where
--     memoize = $(deriveMemoize ''T)
-- @
deriveMemoize  Name  ExpQ
deriveMemoize :: Name -> ExpQ
deriveMemoize name0 :: Name
name0 = do
  (_, _, cons :: [(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' name0 :: Name
name0 mindices :: Maybe [Int]
mindices = do
  (name :: Name
name, tvbs :: [TyVarBndr]
tvbs, cons :: [(Name, Int)]
cons)  Name -> Q (Name, [TyVarBndr], [(Name, Int)])
checkName Name
name0
  let tvs :: [Name]
tvs = [TyVarBndr] -> [Name]
forall a. [a] -> [Name]
freshNames [TyVarBndr]
tvbs
  Dec
inst  CxtQ -> TypeQ -> [DecQ] -> DecQ
instanceD
           (Maybe [Int] -> [TyVarBndr] -> [Name] -> CxtQ
buildContext Maybe [Int]
mindices [TyVarBndr]
tvbs [Name]
tvs)
           (Name -> [Name] -> TypeQ
buildHead Name
name [Name]
tvs)
           [[(Name, Int)] -> DecQ
buildMethodDec [(Name, Int)]
cons]
  [Dec] -> Q [Dec]
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, [TyVarBndr], [(Name, Int)])
checkName :: Name -> Q (Name, [TyVarBndr], [(Name, Int)])
checkName name0 :: Name
name0 = do
  Info
info             Name -> Q Info
reify Name
name0
  case Info
info of
#if MIN_VERSION_template_haskell(2,11,0)
    TyConI (DataD _ name :: Name
name tvbs :: [TyVarBndr]
tvbs _ cons :: [Con]
cons _)
#else
    TyConI (DataD _ name tvbs cons _)
#endif
                (Name, [TyVarBndr], [(Name, Int)])
-> Q (Name, [TyVarBndr], [(Name, Int)])
forall (m :: * -> *) a. Monad m => a -> m a
return (Name
name, [TyVarBndr]
tvbs, Con -> (Name, Int)
stdizeCon (Con -> (Name, Int)) -> [Con] -> [(Name, Int)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Con]
cons)
#if MIN_VERSION_template_haskell(2,11,0)
    TyConI (NewtypeD _ name :: Name
name tvbs :: [TyVarBndr]
tvbs _ con :: Con
con _)
#else
    TyConI (NewtypeD _ name tvbs con _)
#endif
                (Name, [TyVarBndr], [(Name, Int)])
-> Q (Name, [TyVarBndr], [(Name, Int)])
forall (m :: * -> *) a. Monad m => a -> m a
return (Name
name, [TyVarBndr]
tvbs, [Con -> (Name, Int)
stdizeCon Con
con])
    _           String -> Q (Name, [TyVarBndr], [(Name, Int)])
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Q (Name, [TyVarBndr], [(Name, Int)]))
-> String -> Q (Name, [TyVarBndr], [(Name, Int)])
forall a b. (a -> b) -> a -> b
$
      "deriveMemoizable: Can't derive a Memoizable instance for `" String -> String -> String
forall a. [a] -> [a] -> [a]
++
      Name -> String
forall a. Show a => a -> String
show Name
name0 String -> String -> String
forall a. [a] -> [a] -> [a]
++ "' because it isn't a type constructor."
  where
    stdizeCon :: Con -> (Name, Int)
stdizeCon (NormalC name :: Name
name params :: [BangType]
params) = (Name
name, [BangType] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [BangType]
params)
    stdizeCon (RecC name :: Name
name fields :: [VarBangType]
fields)    = (Name
name, [VarBangType] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [VarBangType]
fields)
    stdizeCon (InfixC _ name :: Name
name _)     = (Name
name, 2)
    stdizeCon (ForallC _ _ con :: Con
con)     = Con -> (Name, Int)
stdizeCon Con
con

-- | 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 :: [a] -> [Name]
freshNames xs :: [a]
xs = Int -> [Name] -> [Name]
forall a. Int -> [a] -> [a]
take ([a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs) [Name]
alphabet
  where
  alphabet :: [Name]
alphabet = [ String -> Name
mkName (Char
cChar -> String -> String
forall a. a -> [a] -> [a]
:String
s)
             | String
s  "" String -> [String] -> [String]
forall a. a -> [a] -> [a]
: (Integer -> String
forall a. Show a => a -> String
show (Integer -> String) -> [Integer] -> [String]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [1  Integer ..])
             , Char
c  ['a' .. '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]  [TyVarBndr]  [Name]  CxtQ
buildContext :: Maybe [Int] -> [TyVarBndr] -> [Name] -> CxtQ
buildContext mindices :: Maybe [Int]
mindices tvbs :: [TyVarBndr]
tvbs tvs :: [Name]
tvs =
#if MIN_VERSION_template_haskell(2,10,0)
  [TypeQ] -> CxtQ
cxt (TypeQ -> TypeQ -> TypeQ
appT (Name -> TypeQ
conT ''Memoizable) (TypeQ -> TypeQ) -> (Name -> TypeQ) -> Name -> TypeQ
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> TypeQ
varT (Name -> TypeQ) -> [Name] -> [TypeQ]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Name]
cxttvs)
#else
  cxt (classP ''Memoizable . (:[]) . varT <$> cxttvs)
#endif
  where
  cxttvs :: [Name]
cxttvs = case Maybe [Int]
mindices of
    Just ixs :: [Int]
ixs  (Int -> Bool) -> [Int] -> [Name] -> [Name]
forall a b. (a -> Bool) -> [a] -> [b] -> [b]
filterBy (Int -> [Int] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Int]
ixs) [1 ..] [Name]
tvs
    Nothing   (TyVarBndr -> Bool) -> [TyVarBndr] -> [Name] -> [Name]
forall a b. (a -> Bool) -> [a] -> [b] -> [b]
filterBy TyVarBndr -> Bool
isStar       [TyVarBndr]
tvbs   [Name]
tvs
  --
  isStar :: TyVarBndr -> Bool
isStar (PlainTV _) = Bool
True
#if __GLASGOW_HASKELL__ >= 706
  isStar (KindedTV _ StarT) = Bool
True
#else
  isStar (KindedTV _ StarK) = True
#endif
  isStar (KindedTV _ _) = Bool
False
  --
  filterBy  (a  Bool)  [a]  [b]  [b]
  filterBy :: (a -> Bool) -> [a] -> [b] -> [b]
filterBy p :: a -> Bool
p xs :: [a]
xs ys :: [b]
ys = (a, b) -> b
forall a b. (a, b) -> b
snd ((a, b) -> b) -> [(a, b)] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((a, b) -> Bool) -> [(a, b)] -> [(a, b)]
forall a. (a -> Bool) -> [a] -> [a]
filter (a -> Bool
p (a -> Bool) -> ((a, b) -> a) -> (a, b) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, b) -> a
forall a b. (a, b) -> a
fst) ([a] -> [b] -> [(a, b)]
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 :: [Name]
tvs = 
  TypeQ -> TypeQ -> TypeQ
appT (Name -> TypeQ
conT ''Memoizable) ((TypeQ -> TypeQ -> TypeQ) -> TypeQ -> [TypeQ] -> TypeQ
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl TypeQ -> TypeQ -> TypeQ
appT (Name -> TypeQ
conT Name
name) (Name -> TypeQ
varT (Name -> TypeQ) -> [Name] -> [TypeQ]
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)] -> DecQ
buildMethodDec cons :: [(Name, Int)]
cons = do
  PatQ -> BodyQ -> [DecQ] -> DecQ
valD (Name -> PatQ
varP 'memoize)
    (ExpQ -> BodyQ
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 cons :: [(Name, Int)]
cons = do
  Name
f       String -> Q Name
newName "f"
  Name
look    String -> Q Name
newName "look"
  [Name]
caches  ((Name, Int) -> Q Name) -> [(Name, Int)] -> Q [Name]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\ _ -> String -> Q Name
newName "cache") [(Name, Int)]
cons
  PatQ -> ExpQ -> ExpQ
lam1E (Name -> PatQ
varP Name
f)
    ([DecQ] -> ExpQ -> ExpQ
letE
      (Name -> [(Name, Int)] -> [Name] -> DecQ
buildLookup Name
look [(Name, Int)]
cons [Name]
caches
        DecQ -> [DecQ] -> [DecQ]
forall a. a -> [a] -> [a]
: ((Name, Int) -> Name -> DecQ) -> [(Name, Int)] -> [Name] -> [DecQ]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (Name -> (Name, Int) -> Name -> DecQ
buildCache Name
f) [(Name, Int)]
cons [Name]
caches)
      (Name -> ExpQ
varE Name
look))

-- | Build the look function by building a clause for each constructor
--   of the datatype.
buildLookup  Name  [(Name, Int)]  [Name]  DecQ
buildLookup :: Name -> [(Name, Int)] -> [Name] -> DecQ
buildLookup look :: Name
look cons :: [(Name, Int)]
cons caches :: [Name]
caches =
  Name -> [ClauseQ] -> DecQ
funD Name
look (((Name, Int) -> Name -> ClauseQ)
-> [(Name, Int)] -> [Name] -> [ClauseQ]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (Name, Int) -> Name -> ClauseQ
buildLookupClause [(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.
buildLookupClause  (Name, Int)  Name  ClauseQ
buildLookupClause :: (Name, Int) -> Name -> ClauseQ
buildLookupClause (con :: Name
con, arity :: Int
arity) cache :: Name
cache = do
  [Name]
params  Int -> Q Name -> Q [Name]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
arity (String -> Q Name
newName "a")
  [PatQ] -> BodyQ -> [DecQ] -> ClauseQ
clause [Name -> [PatQ] -> PatQ
conP Name
con (Name -> PatQ
varP (Name -> PatQ) -> [Name] -> [PatQ]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Name]
params)]
         (ExpQ -> BodyQ
normalB ((ExpQ -> ExpQ -> ExpQ) -> ExpQ -> [ExpQ] -> ExpQ
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl ExpQ -> ExpQ -> ExpQ
appE (Name -> ExpQ
varE Name
cache) (Name -> ExpQ
varE (Name -> ExpQ) -> [Name] -> [ExpQ]
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 -> DecQ
buildCache f :: Name
f (con :: Name
con, arity :: Int
arity) cache :: Name
cache =
  PatQ -> BodyQ -> [DecQ] -> DecQ
valD (Name -> PatQ
varP Name
cache) (ExpQ -> BodyQ
normalB (Int -> Name -> ExpQ -> ExpQ
composeMemos Int
arity Name
f (Name -> ExpQ
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 0     f :: Name
f arg :: ExpQ
arg = [| $(varE f) $arg |]
composeMemos arity :: Int
arity f :: Name
f arg :: ExpQ
arg = do
  [| memoize $ \b -> $(composeMemos (arity - 1) f [| $arg b |]) |]