{-# LANGUAGE CPP              #-}
{-# LANGUAGE ConstraintKinds  #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TemplateHaskell  #-}
{-# LANGUAGE TypeOperators    #-}
--------------------------------------------------------------------------------
-- |
-- Module      :  Data.Comp.Derive.HaskellStrict
-- Copyright   :  (c) 2011 Patrick Bahr
-- License     :  BSD3
-- Maintainer  :  Patrick Bahr <paba@diku.dk>
-- Stability   :  experimental
-- Portability :  non-portable (GHC Extensions)
--
-- Automatically derive instances of @HaskellStrict@.
--
--------------------------------------------------------------------------------

module Data.Comp.Derive.HaskellStrict
    (
     makeHaskellStrict
     , haskellStrict
     , haskellStrict'
    ) where

import Control.Monad hiding (mapM, sequence)
import Data.Comp.Derive.Utils
import Data.Comp.Sum
import Data.Comp.Thunk
import Data.Foldable hiding (any, or)
import Data.Maybe
import Data.Traversable
import Language.Haskell.TH
import Prelude hiding (foldl, foldr, mapM, sequence)
import qualified Prelude as P (all, foldl, foldr, mapM)


class HaskellStrict f where
    thunkSequence :: (Monad m) => f (TermT m g) -> m (f (TermT m g))
    thunkSequenceInject :: (Monad m, f :<: m :+: g) => f (TermT m g) -> TermT m g
    thunkSequenceInject f (TermT m g)
t = m (TermT m g) -> TermT m g
forall (m :: * -> *) h (f :: * -> *) a.
m (CxtT m h f a) -> CxtT m h f a
thunk (m (TermT m g) -> TermT m g) -> m (TermT m g) -> TermT m g
forall a b. (a -> b) -> a -> b
$ (f (TermT m g) -> TermT m g) -> m (f (TermT m g)) -> m (TermT m g)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM f (TermT m g) -> TermT m g
forall (g :: * -> *) (f :: * -> *) h a.
(g :<: f) =>
g (Cxt h f a) -> Cxt h f a
inject (m (f (TermT m g)) -> m (TermT m g))
-> m (f (TermT m g)) -> m (TermT m g)
forall a b. (a -> b) -> a -> b
$ f (TermT m g) -> m (f (TermT m g))
forall (f :: * -> *) (m :: * -> *) (g :: * -> *).
(HaskellStrict f, Monad m) =>
f (TermT m g) -> m (f (TermT m g))
thunkSequence f (TermT m g)
t
    thunkSequenceInject' :: (Monad m, f :<: m :+: g) => f (TermT m g) -> TermT m g
    thunkSequenceInject' = f (TermT m g) -> TermT m g
forall (f :: * -> *) (m :: * -> *) (g :: * -> *).
(HaskellStrict f, Monad m, f :<: (m :+: g)) =>
f (TermT m g) -> TermT m g
thunkSequenceInject

haskellStrict :: (Monad m, HaskellStrict f, f :<: m :+: g) => f (TermT m g) -> TermT m g
haskellStrict :: f (TermT m g) -> TermT m g
haskellStrict = f (TermT m g) -> TermT m g
forall (f :: * -> *) (m :: * -> *) (g :: * -> *).
(HaskellStrict f, Monad m, f :<: (m :+: g)) =>
f (TermT m g) -> TermT m g
thunkSequenceInject

haskellStrict' :: (Monad m, HaskellStrict f, f :<: m :+: g) => f (TermT m g) -> TermT m g
haskellStrict' :: f (TermT m g) -> TermT m g
haskellStrict' = f (TermT m g) -> TermT m g
forall (f :: * -> *) (m :: * -> *) (g :: * -> *).
(HaskellStrict f, Monad m, f :<: (m :+: g)) =>
f (TermT m g) -> TermT m g
thunkSequenceInject'

deepThunk :: t -> ExpQ
deepThunk t
d = t -> ExpQ -> ExpQ
forall t. (Eq t, Num t) => t -> ExpQ -> ExpQ
iter t
d [|thunkSequence|]
    where iter :: t -> ExpQ -> ExpQ
iter t
0 ExpQ
_ = [|whnf'|]
          iter t
1 ExpQ
e = ExpQ
e
          iter t
n ExpQ
e = t -> ExpQ -> ExpQ
iter (t
nt -> t -> t
forall a. Num a => a -> a -> a
-t
1) ([|mapM|] ExpQ -> ExpQ -> ExpQ
`appE` ExpQ
e)

{-| Derive an instance of 'HaskellStrict' for a type constructor of any
  first-order kind taking at least one argument. -}
makeHaskellStrict :: Name -> Q [Dec]
makeHaskellStrict :: Name -> Q [Dec]
makeHaskellStrict Name
fname = do
  Just (DataInfo Cxt
_cxt Name
name [TyVarBndr]
args [Con]
constrs [DerivClause]
_deriving) <- Q Info -> Q (Maybe DataInfo)
abstractNewtypeQ (Q Info -> Q (Maybe DataInfo)) -> Q Info -> Q (Maybe DataInfo)
forall a b. (a -> b) -> a -> b
$ Name -> Q Info
reify Name
fname
  let fArg :: Type
fArg = Name -> Type
VarT (Name -> Type) -> (TyVarBndr -> Name) -> TyVarBndr -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TyVarBndr -> Name
tyVarBndrName (TyVarBndr -> Type) -> TyVarBndr -> Type
forall a b. (a -> b) -> a -> b
$ [TyVarBndr] -> TyVarBndr
forall a. [a] -> a
last [TyVarBndr]
args
      argNames :: Cxt
argNames = (TyVarBndr -> Type) -> [TyVarBndr] -> Cxt
forall a b. (a -> b) -> [a] -> [b]
map (Name -> Type
VarT (Name -> Type) -> (TyVarBndr -> Name) -> TyVarBndr -> Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TyVarBndr -> Name
tyVarBndrName) ([TyVarBndr] -> [TyVarBndr]
forall a. [a] -> [a]
init [TyVarBndr]
args)
      complType :: Type
complType = (Type -> Type -> Type) -> Type -> Cxt -> Type
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Type -> Type -> Type
AppT (Name -> Type
ConT Name
name) Cxt
argNames
      classType :: Type
classType = Type -> Type -> Type
AppT (Name -> Type
ConT ''HaskellStrict) Type
complType
  [(Name, [[Int]])]
constrs_ <- (Con -> Q (Name, [[Int]])) -> [Con] -> Q [(Name, [[Int]])]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
P.mapM (((Name, [(Bang, Type)], Maybe Type) -> (Name, [[Int]]))
-> Q (Name, [(Bang, Type)], Maybe Type) -> Q (Name, [[Int]])
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM (Type -> (Name, [(Bang, Type)], Maybe Type) -> (Name, [[Int]])
forall a. Type -> (a, [(Bang, Type)], Maybe Type) -> (a, [[Int]])
isFarg Type
fArg) (Q (Name, [(Bang, Type)], Maybe Type) -> Q (Name, [[Int]]))
-> (Con -> Q (Name, [(Bang, Type)], Maybe Type))
-> Con
-> Q (Name, [[Int]])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Con -> Q (Name, [(Bang, Type)], Maybe Type)
normalConStrExp) [Con]
constrs
  if ((Name, [[Int]]) -> Bool -> Bool)
-> Bool -> [(Name, [[Int]])] -> Bool
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\ (Name, [[Int]])
y Bool
x -> Bool
x Bool -> Bool -> Bool
&& ([Int] -> Bool) -> [[Int]] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
P.all [Int] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ((Name, [[Int]]) -> [[Int]]
forall a b. (a, b) -> b
snd (Name, [[Int]])
y)) Bool
True [(Name, [[Int]])]
constrs_
   then do
     Dec
sequenceDecl <- PatQ -> BodyQ -> [DecQ] -> DecQ
valD (Name -> PatQ
varP 'thunkSequence) (ExpQ -> BodyQ
normalB [|return|]) []
     Dec
injectDecl <- PatQ -> BodyQ -> [DecQ] -> DecQ
valD (Name -> PatQ
varP 'thunkSequenceInject) (ExpQ -> BodyQ
normalB [|inject|]) []
     Dec
injectDecl' <- PatQ -> BodyQ -> [DecQ] -> DecQ
valD (Name -> PatQ
varP 'thunkSequenceInject') (ExpQ -> BodyQ
normalB [|inject|]) []
     [Dec] -> Q [Dec]
forall (m :: * -> *) a. Monad m => a -> m a
return [Cxt -> Type -> [Dec] -> Dec
mkInstanceD [] Type
classType [Dec
sequenceDecl, Dec
injectDecl, Dec
injectDecl']]
   else do
     ([Clause]
sc',[Match]
matchPat,[Clause]
ic') <- ([(Clause, Match, Clause)] -> ([Clause], [Match], [Clause]))
-> Q [(Clause, Match, Clause)] -> Q ([Clause], [Match], [Clause])
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM [(Clause, Match, Clause)] -> ([Clause], [Match], [Clause])
forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 (Q [(Clause, Match, Clause)] -> Q ([Clause], [Match], [Clause]))
-> Q [(Clause, Match, Clause)] -> Q ([Clause], [Match], [Clause])
forall a b. (a -> b) -> a -> b
$ ((Name, [[Int]]) -> Q (Clause, Match, Clause))
-> [(Name, [[Int]])] -> Q [(Clause, Match, Clause)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
P.mapM (Name, [[Int]]) -> Q (Clause, Match, Clause)
forall t.
(Eq t, Num t) =>
(Name, [[t]]) -> Q (Clause, Match, Clause)
mkClauses [(Name, [[Int]])]
constrs_
     Name
xn <- String -> Q Name
newName String
"x"
     Exp
doThunk <- [|thunk|]
     let sequenceDecl :: Dec
sequenceDecl = Name -> [Clause] -> Dec
FunD 'thunkSequence [Clause]
sc'
         injectDecl :: Dec
injectDecl = Name -> [Clause] -> Dec
FunD 'thunkSequenceInject [[Pat] -> Body -> [Dec] -> Clause
Clause [Name -> Pat
VarP Name
xn] (Exp -> Body
NormalB (Exp
doThunk Exp -> Exp -> Exp
`AppE` Exp -> [Match] -> Exp
CaseE (Name -> Exp
VarE Name
xn) [Match]
matchPat)) []]
         injectDecl' :: Dec
injectDecl' = Name -> [Clause] -> Dec
FunD 'thunkSequenceInject' [Clause]
ic'
     [Dec] -> Q [Dec]
forall (m :: * -> *) a. Monad m => a -> m a
return [Cxt -> Type -> [Dec] -> Dec
mkInstanceD [] Type
classType [Dec
sequenceDecl, Dec
injectDecl, Dec
injectDecl']]
      where isFarg :: Type -> (a, [(Bang, Type)], Maybe Type) -> (a, [[Int]])
isFarg Type
fArg (a
constr, [(Bang, Type)]
args, Maybe Type
gadtTy) = (a
constr, ((Bang, Type) -> [Int]) -> [(Bang, Type)] -> [[Int]]
forall a b. (a -> b) -> [a] -> [b]
map (Type -> (Bang, Type) -> [Int]
containsStr (Type -> Maybe Type -> Type
getUnaryFArg Type
fArg Maybe Type
gadtTy)) [(Bang, Type)]
args)
            
#if __GLASGOW_HASKELL__ < 800
            containsStr fArg (IsStrict,ty) = ty `containsType'` fArg
            containsStr fArg (Unpacked,ty) = ty `containsType'` fArg
#else
            containsStr :: Type -> (Bang, Type) -> [Int]
containsStr Type
fArg (Bang SourceUnpackedness
_ SourceStrictness
SourceStrict,Type
ty) = Type
ty Type -> Type -> [Int]
`containsType'` Type
fArg
            containsStr Type
fArg (Bang SourceUnpackedness
SourceUnpack SourceStrictness
_,Type
ty) = Type
ty Type -> Type -> [Int]
`containsType'` Type
fArg
#endif
            containsStr Type
_ (Bang, Type)
_ = []

            filterVar :: (t -> t -> p) -> (t -> p) -> [t] -> t -> p
filterVar t -> t -> p
_ t -> p
nonFarg [] t
x  = t -> p
nonFarg t
x
            filterVar t -> t -> p
farg t -> p
_ [t
depth] t
x = t -> t -> p
farg t
depth t
x
            filterVar t -> t -> p
_ t -> p
_ [t]
_ t
_ = String -> p
forall a. HasCallStack => String -> a
error String
"functor variable occurring twice in argument type"
            filterVars :: [[t]] -> [t] -> (t -> t -> c) -> (t -> c) -> [c]
filterVars [[t]]
args [t]
varNs t -> t -> c
farg t -> c
nonFarg = ([t] -> t -> c) -> [[t]] -> [t] -> [c]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith ((t -> t -> c) -> (t -> c) -> [t] -> t -> c
forall t t p. (t -> t -> p) -> (t -> p) -> [t] -> t -> p
filterVar t -> t -> c
farg t -> c
nonFarg) [[t]]
args [t]
varNs
            mkCPat :: Name -> [Name] -> Pat
mkCPat Name
constr [Name]
varNs = Name -> [Pat] -> Pat
ConP Name
constr ([Pat] -> Pat) -> [Pat] -> Pat
forall a b. (a -> b) -> a -> b
$ (Name -> Pat) -> [Name] -> [Pat]
forall a b. (a -> b) -> [a] -> [b]
map Name -> Pat
mkPat [Name]
varNs
            mkPat :: Name -> Pat
mkPat = Name -> Pat
VarP
            mkClauses :: (Name, [[t]]) -> Q (Clause, Match, Clause)
mkClauses (Name
constr, [[t]]
args) =
                do [Name]
varNs <- Int -> String -> Q [Name]
newNames ([[t]] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[t]]
args) String
"x"
                   let pat :: Pat
pat = Name -> [Name] -> Pat
mkCPat Name
constr [Name]
varNs
                       fvars :: [(t, Name)]
fvars = [Maybe (t, Name)] -> [(t, Name)]
forall a. [Maybe a] -> [a]
catMaybes ([Maybe (t, Name)] -> [(t, Name)])
-> [Maybe (t, Name)] -> [(t, Name)]
forall a b. (a -> b) -> a -> b
$ [[t]]
-> [Name]
-> (t -> Name -> Maybe (t, Name))
-> (Name -> Maybe (t, Name))
-> [Maybe (t, Name)]
forall t t c. [[t]] -> [t] -> (t -> t -> c) -> (t -> c) -> [c]
filterVars [[t]]
args [Name]
varNs (((t, Name) -> Maybe (t, Name)) -> t -> Name -> Maybe (t, Name)
forall a b c. ((a, b) -> c) -> a -> b -> c
curry (t, Name) -> Maybe (t, Name)
forall a. a -> Maybe a
Just) (Maybe (t, Name) -> Name -> Maybe (t, Name)
forall a b. a -> b -> a
const Maybe (t, Name)
forall a. Maybe a
Nothing)
                       allVars :: [ExpQ]
allVars = (Name -> ExpQ) -> [Name] -> [ExpQ]
forall a b. (a -> b) -> [a] -> [b]
map Name -> ExpQ
varE [Name]
varNs
                       conAp :: ExpQ
conAp = (ExpQ -> ExpQ -> ExpQ) -> ExpQ -> [ExpQ] -> ExpQ
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
P.foldl ExpQ -> ExpQ -> ExpQ
appE (Name -> ExpQ
conE Name
constr) [ExpQ]
allVars
                       conBind :: (t, Name) -> ExpQ -> ExpQ
conBind (t
d, Name
x) ExpQ
y = [| $(deepThunk d `appE` varE x)  >>= $(lamE [varP x] y)|]
                   Exp
bodySC' <- ((t, Name) -> ExpQ -> ExpQ) -> ExpQ -> [(t, Name)] -> ExpQ
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
P.foldr (t, Name) -> ExpQ -> ExpQ
forall t. (Eq t, Num t) => (t, Name) -> ExpQ -> ExpQ
conBind [|return $conAp|] [(t, Name)]
fvars
                   let sc' :: Clause
sc' = [Pat] -> Body -> [Dec] -> Clause
Clause [Pat
pat] (Exp -> Body
NormalB Exp
bodySC') []
                   Exp
bodyMatch <- case [(t, Name)]
fvars of
                             [] -> [|return (inject $conAp)|]
                             [(t, Name)]
_ -> ((t, Name) -> ExpQ -> ExpQ) -> ExpQ -> [(t, Name)] -> ExpQ
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
P.foldr (t, Name) -> ExpQ -> ExpQ
forall t. (Eq t, Num t) => (t, Name) -> ExpQ -> ExpQ
conBind [|return (inject $conAp)|] [(t, Name)]
fvars
                   let matchPat :: Match
matchPat = Pat -> Body -> [Dec] -> Match
Match Pat
pat (Exp -> Body
NormalB Exp
bodyMatch) []
                   Exp
bodyIC' <- case [(t, Name)]
fvars of
                             [] -> [|inject $conAp|]
                             [(t, Name)]
_ -> [| thunk |] ExpQ -> ExpQ -> ExpQ
`appE` ((t, Name) -> ExpQ -> ExpQ) -> ExpQ -> [(t, Name)] -> ExpQ
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
P.foldr (t, Name) -> ExpQ -> ExpQ
forall t. (Eq t, Num t) => (t, Name) -> ExpQ -> ExpQ
conBind [|return (inject $conAp)|] [(t, Name)]
fvars
                   let ic' :: Clause
ic' = [Pat] -> Body -> [Dec] -> Clause
Clause [Pat
pat] (Exp -> Body
NormalB Exp
bodyIC') []
                   (Clause, Match, Clause) -> Q (Clause, Match, Clause)
forall (m :: * -> *) a. Monad m => a -> m a
return (Clause
sc', Match
matchPat, Clause
ic')