{-# 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 = forall (m :: * -> *) h (f :: * -> *) a.
m (CxtT m h f a) -> CxtT m h f a
thunk forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM forall (g :: * -> *) (f :: * -> *) h a.
(g :<: f) =>
g (Cxt h f a) -> Cxt h f a
inject forall a b. (a -> b) -> a -> b
$ 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' = 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 :: forall (m :: * -> *) (f :: * -> *) (g :: * -> *).
(Monad m, HaskellStrict f, f :<: (m :+: g)) =>
f (TermT m g) -> TermT m g
haskellStrict = 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' :: forall (m :: * -> *) (f :: * -> *) (g :: * -> *).
(Monad m, HaskellStrict f, f :<: (m :+: g)) =>
f (TermT m g) -> TermT m g
haskellStrict' = forall (f :: * -> *) (m :: * -> *) (g :: * -> *).
(HaskellStrict f, Monad m, f :<: (m :+: g)) =>
f (TermT m g) -> TermT m g
thunkSequenceInject'

deepThunk :: t -> m Exp
deepThunk t
d = forall {t} {m :: * -> *}.
(Eq t, Num t, Quote m) =>
t -> m Exp -> m Exp
iter t
d [|thunkSequence|]
    where iter :: t -> m Exp -> m Exp
iter t
0 m Exp
_ = [|whnf'|]
          iter t
1 m Exp
e = m Exp
e
          iter t
n m Exp
e = t -> m Exp -> m Exp
iter (t
nforall a. Num a => a -> a -> a
-t
1) ([|mapM|] forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
`appE` m Exp
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 flag]
args [Con]
constrs [DerivClause]
_deriving) <- Q Info -> Q (Maybe DataInfo)
abstractNewtypeQ forall a b. (a -> b) -> a -> b
$ Name -> Q Info
reify Name
fname
  let fArg :: Type
fArg = Name -> Type
VarT forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {flag}. TyVarBndr flag -> Name
tyVarBndrName forall a b. (a -> b) -> a -> b
$ forall a. [a] -> a
last [TyVarBndr flag]
args
      argNames :: Cxt
argNames = forall a b. (a -> b) -> [a] -> [b]
map (Name -> Type
VarT forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall {flag}. TyVarBndr flag -> Name
tyVarBndrName) (forall a. [a] -> [a]
init [TyVarBndr flag]
args)
      complType :: Type
complType = 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_ <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
P.mapM (forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM (forall {a}. Type -> (a, [(Bang, Type)], Maybe Type) -> (a, [[Int]])
isFarg Type
fArg) forall b c a. (b -> c) -> (a -> b) -> a -> c
. Con -> Q (Name, [(Bang, Type)], Maybe Type)
normalConStrExp) [Con]
constrs
  if forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\ (Name, [[Int]])
y Bool
x -> Bool
x Bool -> Bool -> Bool
&& forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
P.all forall (t :: * -> *) a. Foldable t => t a -> Bool
null (forall a b. (a, b) -> b
snd (Name, [[Int]])
y)) Bool
True [(Name, [[Int]])]
constrs_
   then do
     Dec
sequenceDecl <- forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Dec
valD (forall (m :: * -> *). Quote m => Name -> m Pat
varP 'thunkSequence) (forall (m :: * -> *). Quote m => m Exp -> m Body
normalB [|return|]) []
     Dec
injectDecl <- forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Dec
valD (forall (m :: * -> *). Quote m => Name -> m Pat
varP 'thunkSequenceInject) (forall (m :: * -> *). Quote m => m Exp -> m Body
normalB [|inject|]) []
     Dec
injectDecl' <- forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Dec
valD (forall (m :: * -> *). Quote m => Name -> m Pat
varP 'thunkSequenceInject') (forall (m :: * -> *). Quote m => m Exp -> m Body
normalB [|inject|]) []
     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') <- forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 forall a b. (a -> b) -> a -> b
$ forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
P.mapM forall {t}.
(Eq t, Num t) =>
(Name, [[t]]) -> Q (Clause, Match, Clause)
mkClauses [(Name, [[Int]])]
constrs_
     Name
xn <- forall (m :: * -> *). Quote m => String -> m 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'
     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, 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 -> t) -> (t -> t) -> [t] -> t -> t
filterVar t -> t -> t
_ t -> t
nonFarg [] t
x  = t -> t
nonFarg t
x
            filterVar t -> t -> t
farg t -> t
_ [t
depth] t
x = t -> t -> t
farg t
depth t
x
            filterVar t -> t -> t
_ t -> t
_ [t]
_ t
_ = forall a. HasCallStack => String -> a
error String
"functor variable occurring twice in argument type"
            filterVars :: [[t]] -> [b] -> (t -> b -> c) -> (b -> c) -> [c]
filterVars [[t]]
args [b]
varNs t -> b -> c
farg b -> c
nonFarg = forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (forall {t} {t} {t}. (t -> t -> t) -> (t -> t) -> [t] -> t -> t
filterVar t -> b -> c
farg b -> c
nonFarg) [[t]]
args [b]
varNs
            mkCPat :: Name -> [Name] -> Pat
mkCPat Name
constr [Name]
varNs = Name -> Cxt -> [Pat] -> Pat
ConP Name
constr [] forall a b. (a -> b) -> a -> b
$ 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 (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 = forall a. [Maybe a] -> [a]
catMaybes forall a b. (a -> b) -> a -> b
$ forall {t} {b} {c}.
[[t]] -> [b] -> (t -> b -> c) -> (b -> c) -> [c]
filterVars [[t]]
args [Name]
varNs (forall a b c. ((a, b) -> c) -> a -> b -> c
curry forall a. a -> Maybe a
Just) (forall a b. a -> b -> a
const forall a. Maybe a
Nothing)
                       allVars :: [Q Exp]
allVars = forall a b. (a -> b) -> [a] -> [b]
map forall (m :: * -> *). Quote m => Name -> m Exp
varE [Name]
varNs
                       conAp :: Q Exp
conAp = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
P.foldl forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
appE (forall (m :: * -> *). Quote m => Name -> m Exp
conE Name
constr) [Q Exp]
allVars
                       conBind :: (t, Name) -> m Exp -> m Exp
conBind (t
d, Name
x) m Exp
y = [| $(deepThunk d `appE` varE x)  >>= $(lamE [varP x] y)|]
                   Exp
bodySC' <- forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
P.foldr forall {m :: * -> *} {t}.
(Quote m, Eq t, Num t) =>
(t, Name) -> m Exp -> m Exp
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)]
_ -> forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
P.foldr forall {m :: * -> *} {t}.
(Quote m, Eq t, Num t) =>
(t, Name) -> m Exp -> m Exp
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 |] forall (m :: * -> *). Quote m => m Exp -> m Exp -> m Exp
`appE` forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
P.foldr forall {m :: * -> *} {t}.
(Quote m, Eq t, Num t) =>
(t, Name) -> m Exp -> m Exp
conBind [|return (inject $conAp)|] [(t, Name)]
fvars
                   let ic' :: Clause
ic' = [Pat] -> Body -> [Dec] -> Clause
Clause [Pat
pat] (Exp -> Body
NormalB Exp
bodyIC') []
                   forall (m :: * -> *) a. Monad m => a -> m a
return (Clause
sc', Match
matchPat, Clause
ic')