{-# LANGUAGE FlexibleContexts #-}

module Language.Haskell.Liquid.Transforms.InlineAux
  ( inlineAux
  )
where
import qualified Language.Haskell.Liquid.UX.Config  as UX
import           Liquid.GHC.API
import           Control.Arrow                  (second)
import qualified Language.Haskell.Liquid.GHC.Misc
                                               as GM
import qualified Data.HashMap.Strict           as M

inlineAux :: UX.Config -> Module -> CoreProgram -> CoreProgram
inlineAux :: Config -> Module -> CoreProgram -> CoreProgram
inlineAux Config
cfg Module
m CoreProgram
cbs =  if Config -> Bool
UX.auxInline Config
cfg then Module
-> (Id -> Bool)
-> (Activation -> Bool)
-> [CoreRule]
-> CoreProgram
-> CoreProgram
occurAnalysePgm Module
m (Bool -> Id -> Bool
forall a b. a -> b -> a
const Bool
False) (Bool -> Activation -> Bool
forall a b. a -> b -> a
const Bool
False) [] ((CoreBind -> CoreBind) -> CoreProgram -> CoreProgram
forall a b. (a -> b) -> [a] -> [b]
map CoreBind -> CoreBind
f CoreProgram
cbs) else CoreProgram
cbs
 where
  f :: CoreBind -> CoreBind
  f :: CoreBind -> CoreBind
f all' :: CoreBind
all'@(NonRec Id
x CoreExpr
e)
    | Just (Id
dfunId, HashMap Id Id
methodToAux) <- Id -> HashMap Id (Id, HashMap Id Id) -> Maybe (Id, HashMap Id Id)
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
M.lookup Id
x HashMap Id (Id, HashMap Id Id)
auxToMethodToAux = Id -> CoreExpr -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec
      Id
x
      (Id -> HashMap Id Id -> CoreExpr -> CoreExpr
inlineAuxExpr Id
dfunId HashMap Id Id
methodToAux CoreExpr
e)
    | Bool
otherwise = CoreBind
all'
  f (Rec [(Id, CoreExpr)]
bs) = [(Id, CoreExpr)] -> CoreBind
forall b. [(b, Expr b)] -> Bind b
Rec (((Id, CoreExpr) -> (Id, CoreExpr))
-> [(Id, CoreExpr)] -> [(Id, CoreExpr)]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Id, CoreExpr) -> (Id, CoreExpr)
g [(Id, CoreExpr)]
bs)
   where
    g :: (Id, CoreExpr) -> (Id, CoreExpr)
g all' :: (Id, CoreExpr)
all'@(Id
x, CoreExpr
e)
      | Just (Id
dfunId, HashMap Id Id
methodToAux) <- Id -> HashMap Id (Id, HashMap Id Id) -> Maybe (Id, HashMap Id Id)
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
M.lookup Id
x HashMap Id (Id, HashMap Id Id)
auxToMethodToAux
      = (Id
x, Id -> HashMap Id Id -> CoreExpr -> CoreExpr
inlineAuxExpr Id
dfunId HashMap Id Id
methodToAux CoreExpr
e)
      | Bool
otherwise
      = (Id, CoreExpr)
all'
  auxToMethodToAux :: HashMap Id (Id, HashMap Id Id)
auxToMethodToAux = [HashMap Id (Id, HashMap Id Id)] -> HashMap Id (Id, HashMap Id Id)
forall a. Monoid a => [a] -> a
mconcat ([HashMap Id (Id, HashMap Id Id)]
 -> HashMap Id (Id, HashMap Id Id))
-> [HashMap Id (Id, HashMap Id Id)]
-> HashMap Id (Id, HashMap Id Id)
forall a b. (a -> b) -> a -> b
$ ((Id, CoreExpr) -> HashMap Id (Id, HashMap Id Id))
-> [(Id, CoreExpr)] -> [HashMap Id (Id, HashMap Id Id)]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Id -> CoreExpr -> HashMap Id (Id, HashMap Id Id))
-> (Id, CoreExpr) -> HashMap Id (Id, HashMap Id Id)
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry Id -> CoreExpr -> HashMap Id (Id, HashMap Id Id)
dfunIdSubst) (CoreProgram -> [(Id, CoreExpr)]
grepDFunIds CoreProgram
cbs)


-- inlineDFun :: DynFlags -> CoreProgram -> IO CoreProgram
-- inlineDFun df cbs = mapM go cbs
--  where
--   go orig@(NonRec x e) | isDFunId x = do
--                            -- e''' <- simplifyExpr df e''
--                            let newBody = mkCoreApps (GM.tracePpr ("substituted type:" ++ GM.showPpr (exprType (mkCoreApps e' (Var <$> binders)))) e') (fmap Var binders)
--                                bind = NonRec (mkWildValBinder (exprType newBody)) newBody
--                            pure $ NonRec x (mkLet bind e)
--                        | otherwise  = pure orig
--    where
--     -- wcBinder = mkWildValBinder t
--     (binders, _) = GM.tracePpr "collectBinders"$ collectBinders e
--     e' = substExprAll empty subst e
--   go recs = pure recs
--   subst = buildDictSubst cbs

-- grab the dictionaries
grepDFunIds :: CoreProgram -> [(DFunId, CoreExpr)]
grepDFunIds :: CoreProgram -> [(Id, CoreExpr)]
grepDFunIds = ((Id, CoreExpr) -> Bool) -> [(Id, CoreExpr)] -> [(Id, CoreExpr)]
forall a. (a -> Bool) -> [a] -> [a]
filter (Id -> Bool
isDFunId (Id -> Bool) -> ((Id, CoreExpr) -> Id) -> (Id, CoreExpr) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Id, CoreExpr) -> Id
forall a b. (a, b) -> a
fst) ([(Id, CoreExpr)] -> [(Id, CoreExpr)])
-> (CoreProgram -> [(Id, CoreExpr)])
-> CoreProgram
-> [(Id, CoreExpr)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CoreProgram -> [(Id, CoreExpr)]
forall b. [Bind b] -> [(b, Expr b)]
flattenBinds

isClassOpAuxOccName :: OccName -> Bool
isClassOpAuxOccName :: OccName -> Bool
isClassOpAuxOccName OccName
occ = case OccName -> [Char]
occNameString OccName
occ of
  Char
'$' : Char
'c' : [Char]
_ -> Bool
True
  [Char]
_             -> Bool
False

isClassOpAuxOf :: Id -> Id -> Bool
isClassOpAuxOf :: Id -> Id -> Bool
isClassOpAuxOf Id
aux Id
method = case OccName -> [Char]
occNameString (OccName -> [Char]) -> OccName -> [Char]
forall a b. (a -> b) -> a -> b
$ Id -> OccName
forall a. NamedThing a => a -> OccName
getOccName Id
aux of
  Char
'$' : Char
'c' : [Char]
rest -> [Char]
rest [Char] -> [Char] -> Bool
forall a. Eq a => a -> a -> Bool
== OccName -> [Char]
occNameString (Id -> OccName
forall a. NamedThing a => a -> OccName
getOccName Id
method)
  [Char]
_                -> Bool
False

dfunIdSubst :: DFunId -> CoreExpr -> M.HashMap Id (Id, M.HashMap Id Id)
dfunIdSubst :: Id -> CoreExpr -> HashMap Id (Id, HashMap Id Id)
dfunIdSubst Id
dfunId CoreExpr
e = [(Id, (Id, HashMap Id Id))] -> HashMap Id (Id, HashMap Id Id)
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
M.fromList ([(Id, (Id, HashMap Id Id))] -> HashMap Id (Id, HashMap Id Id))
-> [(Id, (Id, HashMap Id Id))] -> HashMap Id (Id, HashMap Id Id)
forall a b. (a -> b) -> a -> b
$ [Id] -> [(Id, HashMap Id Id)] -> [(Id, (Id, HashMap Id Id))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Id]
auxIds ((Id, HashMap Id Id) -> [(Id, HashMap Id Id)]
forall a. a -> [a]
repeat (Id
dfunId, HashMap Id Id
methodToAux))
 where
  methodToAux :: HashMap Id Id
methodToAux = [(Id, Id)] -> HashMap Id Id
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
M.fromList
    [ (Id
m, Id
aux) | Id
m <- [Id]
methods, Id
aux <- [Id]
auxIds, Id
aux Id -> Id -> Bool
`isClassOpAuxOf` Id
m ]
  ([Id]
_, [Type]
_, Class
cls, [Type]
_) = Type -> ([Id], [Type], Class, [Type])
tcSplitDFunTy (Id -> Type
idType Id
dfunId)
  auxIds :: [Id]
auxIds = (Id -> Bool) -> [Id] -> [Id]
forall a. (a -> Bool) -> [a] -> [a]
filter (OccName -> Bool
isClassOpAuxOccName (OccName -> Bool) -> (Id -> OccName) -> Id -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Id -> OccName
forall a. NamedThing a => a -> OccName
getOccName) (CoreExpr -> [Id]
exprFreeVarsList CoreExpr
e)
  methods :: [Id]
methods = Class -> [Id]
classAllSelIds Class
cls

inlineAuxExpr :: DFunId -> M.HashMap Id Id -> CoreExpr -> CoreExpr
inlineAuxExpr :: Id -> HashMap Id Id -> CoreExpr -> CoreExpr
inlineAuxExpr Id
dfunId HashMap Id Id
methodToAux = CoreExpr -> CoreExpr
go
 where
  go :: CoreExpr -> CoreExpr
  go :: CoreExpr -> CoreExpr
go (Lam Id
b CoreExpr
body) = Id -> CoreExpr -> CoreExpr
forall b. b -> Expr b -> Expr b
Lam Id
b (CoreExpr -> CoreExpr
go CoreExpr
body)
  go (Let CoreBind
b CoreExpr
body)
    | NonRec Id
x CoreExpr
e <- CoreBind
b, Id -> Bool
isDictId Id
x =
        CoreExpr -> CoreExpr
go (CoreExpr -> CoreExpr) -> CoreExpr -> CoreExpr
forall a b. (a -> b) -> a -> b
$ HasDebugCallStack => Subst -> CoreExpr -> CoreExpr
Subst -> CoreExpr -> CoreExpr
substExpr (Subst -> Id -> CoreExpr -> Subst
extendIdSubst Subst
emptySubst Id
x CoreExpr
e) CoreExpr
body
    | Bool
otherwise = CoreBind -> CoreExpr -> CoreExpr
forall b. Bind b -> Expr b -> Expr b
Let ((CoreExpr -> CoreExpr) -> CoreBind -> CoreBind
forall b. (Expr b -> Expr b) -> Bind b -> Bind b
mapBnd CoreExpr -> CoreExpr
go CoreBind
b) (CoreExpr -> CoreExpr
go CoreExpr
body)
  go (Case CoreExpr
e Id
x Type
t [Alt Id]
alts) = CoreExpr -> Id -> Type -> [Alt Id] -> CoreExpr
forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case (CoreExpr -> CoreExpr
go CoreExpr
e) Id
x Type
t ((Alt Id -> Alt Id) -> [Alt Id] -> [Alt Id]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((CoreExpr -> CoreExpr) -> Alt Id -> Alt Id
forall b. (Expr b -> Expr b) -> Alt b -> Alt b
mapAlt CoreExpr -> CoreExpr
go) [Alt Id]
alts)
  go (Cast CoreExpr
e CoercionR
c       ) = CoreExpr -> CoercionR -> CoreExpr
forall b. Expr b -> CoercionR -> Expr b
Cast (CoreExpr -> CoreExpr
go CoreExpr
e) CoercionR
c
  go (Tick CoreTickish
t CoreExpr
e       ) = CoreTickish -> CoreExpr -> CoreExpr
forall b. CoreTickish -> Expr b -> Expr b
Tick CoreTickish
t (CoreExpr -> CoreExpr
go CoreExpr
e)
  go CoreExpr
e
    | (Var Id
m, [CoreExpr]
args) <- CoreExpr -> (CoreExpr, [CoreExpr])
forall b. Expr b -> (Expr b, [Expr b])
collectArgs CoreExpr
e
    , Just Id
aux <- Id -> HashMap Id Id -> Maybe Id
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
M.lookup Id
m HashMap Id Id
methodToAux
    , CoreExpr
arg : [CoreExpr]
argsNoTy <- (CoreExpr -> Bool) -> [CoreExpr] -> [CoreExpr]
forall a. (a -> Bool) -> [a] -> [a]
dropWhile CoreExpr -> Bool
forall b. Expr b -> Bool
isTypeArg [CoreExpr]
args
    , (Var Id
x, [CoreExpr]
argargs) <- CoreExpr -> (CoreExpr, [CoreExpr])
forall b. Expr b -> (Expr b, [Expr b])
collectArgs CoreExpr
arg
    , Id
x Id -> Id -> Bool
forall a. Eq a => a -> a -> Bool
== Id
dfunId
    = [Char] -> CoreExpr -> CoreExpr
forall a. Outputable a => [Char] -> a -> a
GM.notracePpr ([Char]
"inlining in" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ CoreExpr -> [Char]
forall a. Outputable a => a -> [Char]
GM.showPpr CoreExpr
e)
      (CoreExpr -> CoreExpr) -> CoreExpr -> CoreExpr
forall a b. (a -> b) -> a -> b
$ CoreExpr -> [CoreExpr] -> CoreExpr
mkCoreApps (Id -> CoreExpr
forall b. Id -> Expr b
Var Id
aux) ([CoreExpr]
argargs [CoreExpr] -> [CoreExpr] -> [CoreExpr]
forall a. [a] -> [a] -> [a]
++ (CoreExpr -> CoreExpr
go (CoreExpr -> CoreExpr) -> [CoreExpr] -> [CoreExpr]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [CoreExpr]
argsNoTy))
  go (App CoreExpr
e0 CoreExpr
e1) = CoreExpr -> CoreExpr -> CoreExpr
forall b. Expr b -> Expr b -> Expr b
App (CoreExpr -> CoreExpr
go CoreExpr
e0) (CoreExpr -> CoreExpr
go CoreExpr
e1)
  go CoreExpr
e           = CoreExpr
e


-- modified from Rec.hs
mapBnd :: (Expr b -> Expr b) -> Bind b -> Bind b
mapBnd :: forall b. (Expr b -> Expr b) -> Bind b -> Bind b
mapBnd Expr b -> Expr b
f (NonRec b
b Expr b
e) = b -> Expr b -> Bind b
forall b. b -> Expr b -> Bind b
NonRec b
b (Expr b -> Expr b
f Expr b
e)
mapBnd Expr b -> Expr b
f (Rec [(b, Expr b)]
bs    ) = [(b, Expr b)] -> Bind b
forall b. [(b, Expr b)] -> Bind b
Rec (((b, Expr b) -> (b, Expr b)) -> [(b, Expr b)] -> [(b, Expr b)]
forall a b. (a -> b) -> [a] -> [b]
map ((Expr b -> Expr b) -> (b, Expr b) -> (b, Expr b)
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second Expr b -> Expr b
f) [(b, Expr b)]
bs)

mapAlt :: (Expr b -> Expr b) -> Alt b -> Alt b
mapAlt :: forall b. (Expr b -> Expr b) -> Alt b -> Alt b
mapAlt Expr b -> Expr b
f (Alt AltCon
d [b]
bs Expr b
e) = AltCon -> [b] -> Expr b -> Alt b
forall b. AltCon -> [b] -> Expr b -> Alt b
Alt AltCon
d [b]
bs (Expr b -> Expr b
f Expr b
e)