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


-- | 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.SingleTick
import Rattus.Plugin.CheckSingleTick
import Rattus.Plugin.Utils
import Rattus.Plugin.Annotation

import Prelude hiding ((<>))

import Control.Monad
import Data.Maybe
import Data.Data hiding (tyConName)
import qualified Data.Set as Set

#if __GLASGOW_HASKELL__ >= 900
import GHC.Plugins
import GHC.Tc.Types
#else
import GhcPlugins
import TcRnTypes
#endif

-- | 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
  }


data Options = Options {Options -> Bool
debugMode :: Bool}

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]
opts [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" (Options -> CorePluginPass
strictifyProgram Options :: Bool -> Options
Options{debugMode :: Bool
debugMode = Bool
dmode})
          dmode :: Bool
dmode = CommandLineOption
"debug" CommandLineOption -> [CommandLineOption] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [CommandLineOption]
opts

-- | Apply the following operations to all Rattus definitions in the
-- program:
--
-- * Transform into single tick form (see SingleTick module)
-- * Check whether lazy data types are used (see Strictify module)
-- * Transform into call-by-value form (see Strictify module)

strictifyProgram :: Options -> ModGuts -> CoreM ModGuts
strictifyProgram :: Options -> CorePluginPass
strictifyProgram Options
opts 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 (Options -> ModGuts -> CoreBind -> CoreM CoreBind
strictify Options
opts 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 :: Options -> ModGuts -> CoreBind -> CoreM (CoreBind)
strictify :: Options -> ModGuts -> CoreBind -> CoreM CoreBind
strictify Options
opts 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
                    Expr CoreBndr
e' <- Expr CoreBndr -> CoreM (Expr CoreBndr)
toSingleTick Expr CoreBndr
e
                    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'
                    CheckExpr -> Expr CoreBndr -> CoreM ()
checkExpr CheckExpr :: Set CoreBndr -> Expr CoreBndr -> Bool -> Bool -> CheckExpr
CheckExpr{ recursiveSet :: Set CoreBndr
recursiveSet = [CoreBndr] -> Set CoreBndr
forall a. Ord a => [a] -> Set a
Set.fromList [CoreBndr]
vs, oldExpr :: Expr CoreBndr
oldExpr = Expr CoreBndr
e,
                                         fatalError :: Bool
fatalError = Bool
False, verbose :: Bool
verbose = Options -> Bool
debugMode Options
opts} Expr CoreBndr
e''
                    Expr CoreBndr -> CoreM (Expr CoreBndr)
forall (m :: * -> *) a. Monad m => a -> m a
return 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 Options
opts 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
      -- liftIO $ putStrLn "-------- old --------"
      -- liftIO $ putStrLn (showSDocUnsafe (ppr e))
      Expr CoreBndr
e' <- Expr CoreBndr -> CoreM (Expr CoreBndr)
toSingleTick Expr CoreBndr
e
      -- liftIO $ putStrLn "-------- new --------"
      -- liftIO $ putStrLn (showSDocUnsafe (ppr e'))
      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'
      CheckExpr -> Expr CoreBndr -> CoreM ()
checkExpr CheckExpr :: Set CoreBndr -> Expr CoreBndr -> Bool -> Bool -> CheckExpr
CheckExpr{ recursiveSet :: Set CoreBndr
recursiveSet = Set CoreBndr
forall a. Set a
Set.empty, oldExpr :: Expr CoreBndr
oldExpr = Expr CoreBndr
e,
                           fatalError :: Bool
fatalError = Bool
False, verbose :: Bool
verbose = Options -> Bool
debugMode Options
opts} 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]
  [InternalAnn]
l' <- ModGuts -> CoreBndr -> CoreM [InternalAnn]
forall a. Data a => ModGuts -> CoreBndr -> CoreM [a]
annotationsOn ModGuts
guts CoreBndr
bndr :: CoreM [InternalAnn]
  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) Bool -> Bool -> Bool
&& Bool -> Bool
not (InternalAnn
ExpectError InternalAnn -> [InternalAnn] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [InternalAnn]
l'))

annotationsOn :: (Data a) => ModGuts -> CoreBndr -> CoreM [a]
annotationsOn :: ModGuts -> CoreBndr -> CoreM [a]
annotationsOn ModGuts
guts CoreBndr
bndr = do
#if __GLASGOW_HASKELL__ >= 900
  (_,anns)  <- getAnnotations deserializeWithData guts
  return $
    lookupWithDefaultUFM anns [] (varName bndr) ++
    getModuleAnnotations guts
#else    
  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
#endif