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


-- | The plugin to make it all work.

module AsyncRattus.Plugin (plugin, AsyncRattus(..)) where
import AsyncRattus.Plugin.StableSolver
import AsyncRattus.Plugin.ScopeCheck
import AsyncRattus.Plugin.Strictify
import AsyncRattus.Plugin.SingleTick
import AsyncRattus.Plugin.CheckClockCompatibility
import AsyncRattus.Plugin.Utils
import AsyncRattus.Plugin.Annotation
import AsyncRattus.Plugin.Transform

import Prelude hiding ((<>))

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

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

-- | Use this to enable Asynchronous Rattus' plugin, either by supplying the option
-- @-fplugin=AsyncRattus.Plugin@ directly to GHC, or by including the
-- following pragma in each source file:
-- 
-- > {-# OPTIONS -fplugin=AsyncRattus.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 = case forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find CoreToDo -> Bool
findSamePass [CoreToDo]
todo of       -- check that we don't run the transformation twice
                      Maybe CoreToDo
Nothing -> forall (m :: * -> *) a. Monad m => a -> m a
return (CoreToDo
strPass forall a. a -> [a] -> [a]
: [CoreToDo]
todo) -- (e.g. if the "-fplugin" option is used twice)
                      Maybe CoreToDo
_ -> forall (m :: * -> *) a. Monad m => a -> m a
return [CoreToDo]
todo
    where name :: CommandLineOption
name = CommandLineOption
"Async Rattus strictify"
          strPass :: CoreToDo
strPass = CommandLineOption -> CorePluginPass -> CoreToDo
CoreDoPluginPass CommandLineOption
name (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
          findSamePass :: CoreToDo -> Bool
findSamePass (CoreDoPluginPass CommandLineOption
s CorePluginPass
_) = CommandLineOption
s forall a. Eq a => a -> a -> Bool
== CommandLineOption
name
          findSamePass CoreToDo
_ = Bool
False
          

-- | Apply the following operations to all Asynchronous 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 [(Var, CoreExpr)]
bs) = do
  let debug :: Bool
debug = Options -> Bool
debugMode Options
opts
  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 -> Var -> CoreM Bool
shouldProcessCore ModGuts
guts forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [(Var, CoreExpr)]
bs)
  if Bool
tr then do
    let vs :: [Var]
vs = forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(Var, CoreExpr)]
bs
    [CoreExpr]
es' <- forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\ (Var
v,CoreExpr
e) -> do
      Bool
processCore <- ModGuts -> Var -> CoreM Bool
shouldProcessCore ModGuts
guts Var
v
      if Bool -> Bool
not Bool
processCore
      then do
        forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
debug forall a b. (a -> b) -> a -> b
$ SDoc -> CoreM ()
putMsg forall a b. (a -> b) -> a -> b
$ CommandLineOption -> SDoc
text CommandLineOption
"Skipping binding: " SDoc -> SDoc -> SDoc
<> forall a. Outputable a => a -> SDoc
ppr Var
v
        forall (m :: * -> *) a. Monad m => a -> m a
return CoreExpr
e
      else ModGuts -> Set Var -> Bool -> Var -> CoreExpr -> CoreM CoreExpr
checkAndTransform ModGuts
guts (forall a. Ord a => [a] -> Set a
Set.fromList [Var]
vs) Bool
debug Var
v CoreExpr
e
      ) [(Var, CoreExpr)]
bs
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
debug forall a b. (a -> b) -> a -> b
$ SDoc -> CoreM ()
putMsg forall a b. (a -> b) -> a -> b
$ SDoc
"Plugin | result of transformation: " SDoc -> SDoc -> SDoc
<> forall a. Outputable a => a -> SDoc
ppr [CoreExpr]
es'
    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 [Var]
vs [CoreExpr]
es'))
  else forall (m :: * -> *) a. Monad m => a -> m a
return CoreBind
b
strictify Options
opts ModGuts
guts b :: CoreBind
b@(NonRec Var
v CoreExpr
e) = do
    let debug :: Bool
debug = Options -> Bool
debugMode Options
opts
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
debug forall a b. (a -> b) -> a -> b
$ SDoc -> CoreM ()
putMsg forall a b. (a -> b) -> a -> b
$ CommandLineOption -> SDoc
text CommandLineOption
"Processing binding: " SDoc -> SDoc -> SDoc
<> forall a. Outputable a => a -> SDoc
ppr Var
v SDoc -> SDoc -> SDoc
<> CommandLineOption -> SDoc
text CommandLineOption
" | Non-recursive binding"
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
debug forall a b. (a -> b) -> a -> b
$ SDoc -> CoreM ()
putMsg forall a b. (a -> b) -> a -> b
$ CommandLineOption -> SDoc
text CommandLineOption
"Expr: " SDoc -> SDoc -> SDoc
<> forall a. Outputable a => a -> SDoc
ppr CoreExpr
e
    Bool
processCore <- ModGuts -> Var -> CoreM Bool
shouldProcessCore ModGuts
guts Var
v
    if Bool -> Bool
not Bool
processCore then do
      forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
debug forall a b. (a -> b) -> a -> b
$ SDoc -> CoreM ()
putMsg forall a b. (a -> b) -> a -> b
$ CommandLineOption -> SDoc
text CommandLineOption
"Skipping binding: " SDoc -> SDoc -> SDoc
<> forall a. Outputable a => a -> SDoc
ppr Var
v
      forall (m :: * -> *) a. Monad m => a -> m a
return CoreBind
b
    else do
      CoreExpr
transformed <- ModGuts -> Set Var -> Bool -> Var -> CoreExpr -> CoreM CoreExpr
checkAndTransform ModGuts
guts forall a. Set a
Set.empty Bool
debug Var
v CoreExpr
e
      forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
debug forall a b. (a -> b) -> a -> b
$ SDoc -> CoreM ()
putMsg forall a b. (a -> b) -> a -> b
$ SDoc
"Plugin | result of transformation: " SDoc -> SDoc -> SDoc
<> forall a. Outputable a => a -> SDoc
ppr CoreExpr
transformed
      forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall b. b -> Expr b -> Bind b
NonRec Var
v CoreExpr
transformed

checkAndTransform :: ModGuts -> Set Var -> Bool -> Var -> CoreExpr -> CoreM CoreExpr
checkAndTransform :: ModGuts -> Set Var -> Bool -> Var -> CoreExpr -> CoreM CoreExpr
checkAndTransform ModGuts
guts Set Var
recursiveSet Bool
debug Var
v CoreExpr
e = do
  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
debug forall a b. (a -> b) -> a -> b
$ SDoc -> CoreM ()
putMsg forall a b. (a -> b) -> a -> b
$ CommandLineOption -> SDoc
text CommandLineOption
"Processing binding: " SDoc -> SDoc -> SDoc
<> forall a. Outputable a => a -> SDoc
ppr Var
v
  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
debug forall a b. (a -> b) -> a -> b
$ SDoc -> CoreM ()
putMsg forall a b. (a -> b) -> a -> b
$ CommandLineOption -> SDoc
text CommandLineOption
"Expr: " SDoc -> SDoc -> SDoc
<> forall a. Outputable a => a -> SDoc
ppr CoreExpr
e
  Bool
allowRec <- ModGuts -> Var -> CoreM Bool
allowRecursion ModGuts
guts Var
v
  CoreExpr
singleTick <- CoreExpr -> CoreM CoreExpr
toSingleTick CoreExpr
e
  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
debug forall a b. (a -> b) -> a -> b
$ SDoc -> CoreM ()
putMsg forall a b. (a -> b) -> a -> b
$ CommandLineOption -> SDoc
text CommandLineOption
"Single-tick: " SDoc -> SDoc -> SDoc
<> forall a. Outputable a => a -> SDoc
ppr CoreExpr
singleTick
  Bool
lazy <- ModGuts -> Var -> CoreM Bool
allowLazyData ModGuts
guts Var
v
  CoreExpr
strict <- SCxt -> CoreExpr -> CoreM CoreExpr
strictifyExpr (SrcSpan -> Bool -> SCxt
SCxt (Name -> SrcSpan
nameSrcSpan forall a b. (a -> b) -> a -> b
$ forall a. NamedThing a => a -> Name
getName Var
v) (Bool -> Bool
not Bool
lazy)) CoreExpr
singleTick
  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
debug forall a b. (a -> b) -> a -> b
$ SDoc -> CoreM ()
putMsg forall a b. (a -> b) -> a -> b
$ CommandLineOption -> SDoc
text CommandLineOption
"Strict single-tick: " SDoc -> SDoc -> SDoc
<> forall a. Outputable a => a -> SDoc
ppr CoreExpr
strict
  CheckExpr -> CoreExpr -> CoreM ()
checkExpr CheckExpr{ recursiveSet :: Set Var
recursiveSet = Set Var
recursiveSet, oldExpr :: CoreExpr
oldExpr = CoreExpr
e,
                        verbose :: Bool
verbose = Bool
debug,
                        allowRecExp :: Bool
allowRecExp = Bool
allowRec} CoreExpr
strict
  CoreExpr -> CoreM CoreExpr
transform CoreExpr
strict

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 -> Var -> CoreM Bool
allowLazyData ModGuts
guts Var
bndr = do
  [AsyncRattus]
l <- forall a. Data a => ModGuts -> Var -> CoreM [a]
annotationsOn ModGuts
guts Var
bndr :: CoreM [AsyncRattus]
  forall (m :: * -> *) a. Monad m => a -> m a
return (AsyncRattus
AllowLazyData forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [AsyncRattus]
l)

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

expectError :: ModGuts -> CoreBndr -> CoreM Bool
expectError :: ModGuts -> Var -> CoreM Bool
expectError ModGuts
guts Var
bndr = do
  [InternalAnn]
l <- forall a. Data a => ModGuts -> Var -> CoreM [a]
annotationsOn ModGuts
guts Var
bndr :: CoreM [InternalAnn]
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ InternalAnn
ExpectError forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [InternalAnn]
l


shouldProcessCore :: ModGuts -> CoreBndr -> CoreM Bool
shouldProcessCore :: ModGuts -> Var -> CoreM Bool
shouldProcessCore ModGuts
guts Var
bndr = do
  [AsyncRattus]
l <- forall a. Data a => ModGuts -> Var -> CoreM [a]
annotationsOn ModGuts
guts Var
bndr :: CoreM [AsyncRattus]
  Bool
expectScopeError <- ModGuts -> Var -> CoreM Bool
expectError ModGuts
guts Var
bndr
  forall (m :: * -> *) a. Monad m => a -> m a
return (AsyncRattus
AsyncRattus forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [AsyncRattus]
l Bool -> Bool -> Bool
&& forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
notElem AsyncRattus
NotAsyncRattus [AsyncRattus]
l Bool -> Bool -> Bool
&& Var -> Bool
userFunction Var
bndr Bool -> Bool -> Bool
&& Bool -> Bool
not Bool
expectScopeError)

annotationsOn :: (Data a) => ModGuts -> CoreBndr -> CoreM [a]
annotationsOn :: forall a. Data a => ModGuts -> Var -> CoreM [a]
annotationsOn ModGuts
guts Var
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 [] (Var -> Name
varName Var
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