{-# 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 forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> forall (m :: * -> *) a. Monad m => a -> m a
return TcGblEnv
env

install :: [CommandLineOption] -> [CoreToDo] -> CoreM [CoreToDo]
install :: CorePlugin
install [CommandLineOption]
opts [CoreToDo]
todo = forall (m :: * -> *) a. Monad m => a -> m a
return (CoreToDo
strPass forall a. a -> [a] -> [a]
: [CoreToDo]
todo)
    where strPass :: CoreToDo
strPass = CommandLineOption -> CorePluginPass -> CoreToDo
CoreDoPluginPass CommandLineOption
"Rattus strictify" (Options -> CorePluginPass
strictifyProgram Options{debugMode :: Bool
debugMode = Bool
dmode})
          dmode :: Bool
dmode = CommandLineOption
"debug" 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 <- 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)
  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 <- forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM forall (t :: * -> *). Foldable t => t Bool -> Bool
or (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 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [(CoreBndr, Expr CoreBndr)]
bs)
  if Bool
tr then do
    let vs :: [CoreBndr]
vs = forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(CoreBndr, Expr CoreBndr)]
bs
    [Expr CoreBndr]
es' <- 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
                    Bool
allowRec <- ModGuts -> CoreBndr -> CoreM Bool
allowRecursion ModGuts
guts CoreBndr
v
                    Expr CoreBndr
e'' <- SCxt -> Expr CoreBndr -> CoreM (Expr CoreBndr)
strictifyExpr (SrcSpan -> Bool -> SCxt
SCxt (Name -> SrcSpan
nameSrcSpan forall a b. (a -> b) -> a -> b
$ forall a. NamedThing a => a -> Name
getName CoreBndr
v) (Bool -> Bool
not Bool
lazy))Expr CoreBndr
e'
                    CheckExpr -> Expr CoreBndr -> CoreM ()
checkExpr CheckExpr{ recursiveSet :: Set CoreBndr
recursiveSet = 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,
                                         allowRecExp :: Bool
allowRecExp = Bool
allowRec} Expr CoreBndr
e''
                    forall (m :: * -> *) a. Monad m => a -> m a
return Expr CoreBndr
e'') [(CoreBndr, Expr CoreBndr)]
bs
    forall (m :: * -> *) a. Monad m => a -> m a
return (forall b. [(b, Expr b)] -> Bind b
Rec (forall a b. [a] -> [b] -> [(a, b)]
zip [CoreBndr]
vs [Expr CoreBndr]
es'))
  else 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
      Bool
allowRec <- ModGuts -> CoreBndr -> CoreM Bool
allowRecursion ModGuts
guts CoreBndr
v
      Expr CoreBndr
e'' <- SCxt -> Expr CoreBndr -> CoreM (Expr CoreBndr)
strictifyExpr (SrcSpan -> Bool -> SCxt
SCxt (Name -> SrcSpan
nameSrcSpan forall a b. (a -> b) -> a -> b
$ forall a. NamedThing a => a -> Name
getName CoreBndr
v) (Bool -> Bool
not Bool
lazy)) Expr CoreBndr
e'
      CheckExpr -> Expr CoreBndr -> CoreM ()
checkExpr CheckExpr{ recursiveSet :: Set CoreBndr
recursiveSet = 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,
                           allowRecExp :: Bool
allowRecExp = Bool
allowRec } Expr CoreBndr
e''
      forall (m :: * -> *) a. Monad m => a -> m a
return (forall b. b -> Expr b -> Bind b
NonRec CoreBndr
v Expr CoreBndr
e'')
    else forall (m :: * -> *) a. Monad m => a -> m a
return CoreBind
b

getModuleAnnotations :: Data a => ModGuts -> [a]
getModuleAnnotations :: forall a. Data a => ModGuts -> [a]
getModuleAnnotations ModGuts
guts = [a]
anns'
  where anns :: [Annotation]
anns = forall a. (a -> Bool) -> [a] -> [a]
filter (\Annotation
a-> case Annotation -> CoreAnnTarget
ann_target Annotation
a of
                         ModuleTarget Module
m -> Module
m 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' = forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (forall a. Typeable a => ([Word8] -> a) -> Serialized -> Maybe a
fromSerialized forall a. Data a => [Word8] -> a
deserializeWithData 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 <- forall a. Data a => ModGuts -> CoreBndr -> CoreM [a]
annotationsOn ModGuts
guts CoreBndr
bndr :: CoreM [Rattus]
  forall (m :: * -> *) a. Monad m => a -> m a
return (Rattus
AllowLazyData forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Rattus]
l)

allowRecursion :: ModGuts -> CoreBndr -> CoreM Bool
allowRecursion :: ModGuts -> CoreBndr -> CoreM Bool
allowRecursion ModGuts
guts CoreBndr
bndr = do
  [Rattus]
l <- forall a. Data a => ModGuts -> CoreBndr -> CoreM [a]
annotationsOn ModGuts
guts CoreBndr
bndr :: CoreM [Rattus]
  forall (m :: * -> *) a. Monad m => a -> m a
return (Rattus
AllowRecursion 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 <- forall a. Data a => ModGuts -> CoreBndr -> CoreM [a]
annotationsOn ModGuts
guts CoreBndr
bndr :: CoreM [Rattus]
  [InternalAnn]
l' <- forall a. Data a => ModGuts -> CoreBndr -> CoreM [a]
annotationsOn ModGuts
guts CoreBndr
bndr :: CoreM [InternalAnn]
  forall (m :: * -> *) a. Monad m => a -> m a
return ((Rattus
Rattus forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Rattus]
l Bool -> Bool -> Bool
&& Bool -> Bool
not (Rattus
NotRattus 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 forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [InternalAnn]
l'))

annotationsOn :: (Data a) => ModGuts -> CoreBndr -> CoreM [a]
annotationsOn :: forall a. Data a => ModGuts -> CoreBndr -> CoreM [a]
annotationsOn ModGuts
guts CoreBndr
bndr = do
#if __GLASGOW_HASKELL__ >= 900
  (ModuleEnv [a]
_,NameEnv [a]
anns)  <- forall a.
Typeable a =>
([Word8] -> a) -> ModGuts -> CoreM (ModuleEnv [a], NameEnv [a])
getAnnotations forall a. Data a => [Word8] -> a
deserializeWithData ModGuts
guts
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
    forall key elt.
Uniquable key =>
UniqFM key elt -> elt -> key -> elt
lookupWithDefaultUFM NameEnv [a]
anns [] (CoreBndr -> Name
varName CoreBndr
bndr) forall a. [a] -> [a] -> [a]
++
    forall a. Data a => ModGuts -> [a]
getModuleAnnotations ModGuts
guts
#else    
  anns <- getAnnotations deserializeWithData guts
  return $
    lookupWithDefaultUFM anns [] (varUnique bndr) ++
    getModuleAnnotations guts
#endif