{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveDataTypeable #-}


-- | The plugin to make it all work.

module Rattus.Plugin (plugin, Rattus(..)) where
import Rattus.Plugin.StableSolver
import Rattus.Plugin.ScopeCheck
import Rattus.Plugin.Strictify
import Rattus.Plugin.Utils
import Rattus.Plugin.Annotation

import Prelude hiding ((<>))
import GhcPlugins
import TcRnTypes


import Control.Monad
import Data.Maybe
import Data.Data hiding (tyConName)




-- | Use this to enable Rattus' plugin, either by supplying the option
-- @-fplugin=Rattus.Plugin@ directly to GHC. or by including the
-- following pragma in each source file:
-- 
-- > {-# OPTIONS -fplugin=Rattus.Plugin #-}
plugin :: Plugin
plugin :: Plugin
plugin = Plugin
defaultPlugin {
  installCoreToDos :: CorePlugin
installCoreToDos = CorePlugin
install,
  pluginRecompile :: [CommandLineOption] -> IO PluginRecompile
pluginRecompile = [CommandLineOption] -> IO PluginRecompile
purePlugin,
  typeCheckResultAction :: [CommandLineOption] -> ModSummary -> TcGblEnv -> TcM TcGblEnv
typeCheckResultAction = [CommandLineOption] -> ModSummary -> TcGblEnv -> TcM TcGblEnv
typechecked,
  tcPlugin :: TcPlugin
tcPlugin = TcPlugin
tcStable
  }


typechecked :: [CommandLineOption] -> ModSummary -> TcGblEnv -> TcM TcGblEnv
typechecked :: [CommandLineOption] -> ModSummary -> TcGblEnv -> TcM TcGblEnv
typechecked [CommandLineOption]
_ ModSummary
_ TcGblEnv
env = TcGblEnv -> TcM ()
checkAll TcGblEnv
env TcM () -> TcM TcGblEnv -> TcM TcGblEnv
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> TcGblEnv -> TcM TcGblEnv
forall (m :: * -> *) a. Monad m => a -> m a
return TcGblEnv
env

install :: [CommandLineOption] -> [CoreToDo] -> CoreM [CoreToDo]
install :: CorePlugin
install [CommandLineOption]
_ [CoreToDo]
todo = [CoreToDo] -> CoreM [CoreToDo]
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreToDo
strPass CoreToDo -> [CoreToDo] -> [CoreToDo]
forall a. a -> [a] -> [a]
: [CoreToDo]
todo)
    where strPass :: CoreToDo
strPass = CommandLineOption -> CorePluginPass -> CoreToDo
CoreDoPluginPass CommandLineOption
"Rattus strictify" CorePluginPass
strictifyProgram

strictifyProgram :: ModGuts -> CoreM ModGuts
strictifyProgram :: CorePluginPass
strictifyProgram ModGuts
guts = do
  [CoreBind]
newBinds <- (CoreBind -> CoreM CoreBind) -> [CoreBind] -> CoreM [CoreBind]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (ModGuts -> CoreBind -> CoreM CoreBind
strictify ModGuts
guts) (ModGuts -> [CoreBind]
mg_binds ModGuts
guts)
  CorePluginPass
forall (m :: * -> *) a. Monad m => a -> m a
return ModGuts
guts { mg_binds :: [CoreBind]
mg_binds = [CoreBind]
newBinds }

strictify :: ModGuts -> CoreBind -> CoreM (CoreBind)
strictify :: ModGuts -> CoreBind -> CoreM CoreBind
strictify ModGuts
guts b :: CoreBind
b@(Rec [(CoreBndr, Expr CoreBndr)]
bs) = do
  Bool
tr <- ([Bool] -> Bool) -> CoreM [Bool] -> CoreM Bool
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
or (((CoreBndr, Expr CoreBndr) -> CoreM Bool)
-> [(CoreBndr, Expr CoreBndr)] -> CoreM [Bool]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (ModGuts -> CoreBndr -> CoreM Bool
shouldTransform ModGuts
guts (CoreBndr -> CoreM Bool)
-> ((CoreBndr, Expr CoreBndr) -> CoreBndr)
-> (CoreBndr, Expr CoreBndr)
-> CoreM Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (CoreBndr, Expr CoreBndr) -> CoreBndr
forall a b. (a, b) -> a
fst) [(CoreBndr, Expr CoreBndr)]
bs)
  if Bool
tr then do
    let vs :: [CoreBndr]
vs = ((CoreBndr, Expr CoreBndr) -> CoreBndr)
-> [(CoreBndr, Expr CoreBndr)] -> [CoreBndr]
forall a b. (a -> b) -> [a] -> [b]
map (CoreBndr, Expr CoreBndr) -> CoreBndr
forall a b. (a, b) -> a
fst [(CoreBndr, Expr CoreBndr)]
bs
    [Expr CoreBndr]
es' <- ((CoreBndr, Expr CoreBndr) -> CoreM (Expr CoreBndr))
-> [(CoreBndr, Expr CoreBndr)] -> CoreM [Expr CoreBndr]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\ (CoreBndr
v,Expr CoreBndr
e) -> do
                    Bool
lazy <- ModGuts -> CoreBndr -> CoreM Bool
allowLazyData ModGuts
guts CoreBndr
v
                    SCxt -> Expr CoreBndr -> CoreM (Expr CoreBndr)
strictifyExpr (SrcSpan -> Bool -> SCxt
SCxt (Name -> SrcSpan
nameSrcSpan (Name -> SrcSpan) -> Name -> SrcSpan
forall a b. (a -> b) -> a -> b
$ CoreBndr -> Name
forall a. NamedThing a => a -> Name
getName CoreBndr
v) (Bool -> Bool
not Bool
lazy))Expr CoreBndr
e) [(CoreBndr, Expr CoreBndr)]
bs
    CoreBind -> CoreM CoreBind
forall (m :: * -> *) a. Monad m => a -> m a
return ([(CoreBndr, Expr CoreBndr)] -> CoreBind
forall b. [(b, Expr b)] -> Bind b
Rec ([CoreBndr] -> [Expr CoreBndr] -> [(CoreBndr, Expr CoreBndr)]
forall a b. [a] -> [b] -> [(a, b)]
zip [CoreBndr]
vs [Expr CoreBndr]
es'))
  else CoreBind -> CoreM CoreBind
forall (m :: * -> *) a. Monad m => a -> m a
return CoreBind
b
strictify ModGuts
guts b :: CoreBind
b@(NonRec CoreBndr
v Expr CoreBndr
e) = do
    Bool
tr <- ModGuts -> CoreBndr -> CoreM Bool
shouldTransform ModGuts
guts CoreBndr
v
    if Bool
tr then do
      Bool
lazy <- ModGuts -> CoreBndr -> CoreM Bool
allowLazyData ModGuts
guts CoreBndr
v
      Expr CoreBndr
e' <- SCxt -> Expr CoreBndr -> CoreM (Expr CoreBndr)
strictifyExpr (SrcSpan -> Bool -> SCxt
SCxt (Name -> SrcSpan
nameSrcSpan (Name -> SrcSpan) -> Name -> SrcSpan
forall a b. (a -> b) -> a -> b
$ CoreBndr -> Name
forall a. NamedThing a => a -> Name
getName CoreBndr
v) (Bool -> Bool
not Bool
lazy)) Expr CoreBndr
e
      CoreBind -> CoreM CoreBind
forall (m :: * -> *) a. Monad m => a -> m a
return (CoreBndr -> Expr CoreBndr -> CoreBind
forall b. b -> Expr b -> Bind b
NonRec CoreBndr
v Expr CoreBndr
e')
    else CoreBind -> CoreM CoreBind
forall (m :: * -> *) a. Monad m => a -> m a
return CoreBind
b

getModuleAnnotations :: Data a => ModGuts -> [a]
getModuleAnnotations :: ModGuts -> [a]
getModuleAnnotations ModGuts
guts = [a]
anns'
  where anns :: [Annotation]
anns = (Annotation -> Bool) -> [Annotation] -> [Annotation]
forall a. (a -> Bool) -> [a] -> [a]
filter (\Annotation
a-> case Annotation -> CoreAnnTarget
ann_target Annotation
a of
                         ModuleTarget Module
m -> Module
m Module -> Module -> Bool
forall a. Eq a => a -> a -> Bool
== (ModGuts -> Module
mg_module ModGuts
guts)
                         CoreAnnTarget
_ -> Bool
False) (ModGuts -> [Annotation]
mg_anns ModGuts
guts)
        anns' :: [a]
anns' = (Annotation -> Maybe a) -> [Annotation] -> [a]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (([Word8] -> a) -> Serialized -> Maybe a
forall a. Typeable a => ([Word8] -> a) -> Serialized -> Maybe a
fromSerialized [Word8] -> a
forall a. Data a => [Word8] -> a
deserializeWithData (Serialized -> Maybe a)
-> (Annotation -> Serialized) -> Annotation -> Maybe a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Annotation -> Serialized
ann_value) [Annotation]
anns
  



allowLazyData :: ModGuts -> CoreBndr -> CoreM Bool
allowLazyData :: ModGuts -> CoreBndr -> CoreM Bool
allowLazyData ModGuts
guts CoreBndr
bndr = do
  [Rattus]
l <- ModGuts -> CoreBndr -> CoreM [Rattus]
forall a. Data a => ModGuts -> CoreBndr -> CoreM [a]
annotationsOn ModGuts
guts CoreBndr
bndr :: CoreM [Rattus]
  Bool -> CoreM Bool
forall (m :: * -> *) a. Monad m => a -> m a
return (Rattus
AllowLazyData Rattus -> [Rattus] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Rattus]
l)


shouldTransform :: ModGuts -> CoreBndr -> CoreM Bool
shouldTransform :: ModGuts -> CoreBndr -> CoreM Bool
shouldTransform ModGuts
guts CoreBndr
bndr = do
  [Rattus]
l <- ModGuts -> CoreBndr -> CoreM [Rattus]
forall a. Data a => ModGuts -> CoreBndr -> CoreM [a]
annotationsOn ModGuts
guts CoreBndr
bndr :: CoreM [Rattus]
  Bool -> CoreM Bool
forall (m :: * -> *) a. Monad m => a -> m a
return (Rattus
Rattus Rattus -> [Rattus] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Rattus]
l Bool -> Bool -> Bool
&& Bool -> Bool
not (Rattus
NotRattus Rattus -> [Rattus] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Rattus]
l) Bool -> Bool -> Bool
&& CoreBndr -> Bool
userFunction CoreBndr
bndr)

annotationsOn :: (Data a) => ModGuts -> CoreBndr -> CoreM [a]
annotationsOn :: ModGuts -> CoreBndr -> CoreM [a]
annotationsOn ModGuts
guts CoreBndr
bndr = do
  UniqFM [a]
anns <- ([Word8] -> a) -> ModGuts -> CoreM (UniqFM [a])
forall a.
Typeable a =>
([Word8] -> a) -> ModGuts -> CoreM (UniqFM [a])
getAnnotations [Word8] -> a
forall a. Data a => [Word8] -> a
deserializeWithData ModGuts
guts
  [a] -> CoreM [a]
forall (m :: * -> *) a. Monad m => a -> m a
return ([a] -> CoreM [a]) -> [a] -> CoreM [a]
forall a b. (a -> b) -> a -> b
$
    UniqFM [a] -> [a] -> Unique -> [a]
forall key elt. Uniquable key => UniqFM elt -> elt -> key -> elt
lookupWithDefaultUFM UniqFM [a]
anns [] (CoreBndr -> Unique
varUnique CoreBndr
bndr) [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++
    ModGuts -> [a]
forall a. Data a => ModGuts -> [a]
getModuleAnnotations ModGuts
guts