{-# LANGUAGE OverloadedStrings #-}
module AsyncRattus.Plugin.PrimExpr (
Prim (..),
PrimInfo (..),
function,
prim,
isPrimExpr
) where
import Data.Map (Map)
import qualified Data.Map as Map
import GHC.Plugins
import AsyncRattus.Plugin.Utils
import Prelude hiding ((<>))
data Prim = Delay | Adv | Box | Select
data PrimInfo = DelayApp Var Type | AdvApp Var TypedArg | BoxApp Var | SelectApp Var TypedArg TypedArg
type TypedArg = (Var, Type)
data PartialPrimInfo = PartialPrimInfo {
PartialPrimInfo -> Prim
primPart :: Prim,
PartialPrimInfo -> Var
functionPart :: Var,
PartialPrimInfo -> [Var]
args :: [Var],
PartialPrimInfo -> [Type]
typeArgs :: [Type]
}
instance Outputable PartialPrimInfo where
ppr :: PartialPrimInfo -> SDoc
ppr (PartialPrimInfo Prim
Delay Var
f [Var]
_ [Type]
typeArgs) = String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"PartialPrimInfo { prim = Delay, function = " SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<> Var -> SDoc
forall a. Outputable a => a -> SDoc
ppr Var
f SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<> String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"args = (not printing since it should be undefined) , typeArgs = " SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<> [Type] -> SDoc
forall a. Outputable a => a -> SDoc
ppr [Type]
typeArgs
ppr (PartialPrimInfo Prim
prim Var
f [Var]
args [Type]
typeArgs) = String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"PartialPrimInfo { prim = " SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<> Prim -> SDoc
forall a. Outputable a => a -> SDoc
ppr Prim
prim SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<> String -> SDoc
forall doc. IsLine doc => String -> doc
text String
", function = " SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<> Var -> SDoc
forall a. Outputable a => a -> SDoc
ppr Var
f SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<> String -> SDoc
forall doc. IsLine doc => String -> doc
text String
", args = " SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<> [Var] -> SDoc
forall a. Outputable a => a -> SDoc
ppr [Var]
args SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<> String -> SDoc
forall doc. IsLine doc => String -> doc
text String
", typeArgs = " SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<> [Type] -> SDoc
forall a. Outputable a => a -> SDoc
ppr [Type]
typeArgs
instance Outputable Prim where
ppr :: Prim -> SDoc
ppr Prim
Delay = SDoc
"delay"
ppr Prim
Adv = SDoc
"adv"
ppr Prim
Select = SDoc
"select"
ppr Prim
Box = SDoc
"box"
instance Outputable PrimInfo where
ppr :: PrimInfo -> SDoc
ppr (DelayApp Var
f Type
_) = String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"DelayApp - function " SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<> Var -> SDoc
forall a. Outputable a => a -> SDoc
ppr Var
f
ppr (BoxApp Var
f) = String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"BoxApp - function " SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<> Var -> SDoc
forall a. Outputable a => a -> SDoc
ppr Var
f
ppr (AdvApp Var
f TypedArg
arg) = String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"AdvApp - function " SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<> Var -> SDoc
forall a. Outputable a => a -> SDoc
ppr Var
f SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<> String -> SDoc
forall doc. IsLine doc => String -> doc
text String
" | arg " SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<> TypedArg -> SDoc
forall a. Outputable a => a -> SDoc
ppr TypedArg
arg
ppr (SelectApp Var
f TypedArg
arg TypedArg
arg2) = String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"SelectApp - function " SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<> Var -> SDoc
forall a. Outputable a => a -> SDoc
ppr Var
f SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<> String -> SDoc
forall doc. IsLine doc => String -> doc
text String
" | arg " SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<> TypedArg -> SDoc
forall a. Outputable a => a -> SDoc
ppr TypedArg
arg SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<> String -> SDoc
forall doc. IsLine doc => String -> doc
text String
" | arg2 " SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<> TypedArg -> SDoc
forall a. Outputable a => a -> SDoc
ppr TypedArg
arg2
primMap :: Map FastString Prim
primMap :: Map FastString Prim
primMap = [(FastString, Prim)] -> Map FastString Prim
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList
[(FastString
"delay", Prim
Delay),
(FastString
"adv", Prim
Adv),
(FastString
"select", Prim
Select),
(FastString
"box", Prim
Box)
]
isPrim :: Var -> Maybe Prim
isPrim :: Var -> Maybe Prim
isPrim Var
v = case Var -> Maybe (FastString, FastString)
forall a. NamedThing a => a -> Maybe (FastString, FastString)
getNameModule Var
v of
Just (FastString
name, FastString
mod) | FastString -> Bool
isRattModule FastString
mod -> FastString -> Map FastString Prim -> Maybe Prim
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup FastString
name Map FastString Prim
primMap
Maybe (FastString, FastString)
_ -> Maybe Prim
forall a. Maybe a
Nothing
createPartialPrimInfo :: Prim -> Var -> PartialPrimInfo
createPartialPrimInfo :: Prim -> Var -> PartialPrimInfo
createPartialPrimInfo Prim
prim Var
function =
PartialPrimInfo {
primPart :: Prim
primPart = Prim
prim,
functionPart :: Var
functionPart = Var
function,
args :: [Var]
args = [],
typeArgs :: [Type]
typeArgs = []
}
function :: PrimInfo -> Var
function :: PrimInfo -> Var
function (DelayApp Var
f Type
_) = Var
f
function (BoxApp Var
f) = Var
f
function (AdvApp Var
f TypedArg
_) = Var
f
function (SelectApp Var
f TypedArg
_ TypedArg
_) = Var
f
prim :: PrimInfo -> Prim
prim :: PrimInfo -> Prim
prim (DelayApp {}) = Prim
Delay
prim (BoxApp Var
_) = Prim
Box
prim (AdvApp {}) = Prim
Adv
prim (SelectApp {}) = Prim
Select
validatePartialPrimInfo :: PartialPrimInfo -> Maybe PrimInfo
validatePartialPrimInfo :: PartialPrimInfo -> Maybe PrimInfo
validatePartialPrimInfo (PartialPrimInfo Prim
Select Var
f [Var
arg2V, Var
argV] [Type
arg2T, Type
argT]) = PrimInfo -> Maybe PrimInfo
forall a. a -> Maybe a
Just (PrimInfo -> Maybe PrimInfo) -> PrimInfo -> Maybe PrimInfo
forall a b. (a -> b) -> a -> b
$ Var -> TypedArg -> TypedArg -> PrimInfo
SelectApp Var
f (Var
argV, Type
argT) (Var
arg2V, Type
arg2T)
validatePartialPrimInfo (PartialPrimInfo Prim
Delay Var
f [Var
_] [Type
argT]) = PrimInfo -> Maybe PrimInfo
forall a. a -> Maybe a
Just (PrimInfo -> Maybe PrimInfo) -> PrimInfo -> Maybe PrimInfo
forall a b. (a -> b) -> a -> b
$ Var -> Type -> PrimInfo
DelayApp Var
f Type
argT
validatePartialPrimInfo (PartialPrimInfo {primPart :: PartialPrimInfo -> Prim
primPart = Prim
Box, functionPart :: PartialPrimInfo -> Var
functionPart = Var
f}) = PrimInfo -> Maybe PrimInfo
forall a. a -> Maybe a
Just (PrimInfo -> Maybe PrimInfo) -> PrimInfo -> Maybe PrimInfo
forall a b. (a -> b) -> a -> b
$ Var -> PrimInfo
BoxApp Var
f
validatePartialPrimInfo (PartialPrimInfo Prim
Adv Var
f [Var
argV] [Type
argT]) = PrimInfo -> Maybe PrimInfo
forall a. a -> Maybe a
Just (PrimInfo -> Maybe PrimInfo) -> PrimInfo -> Maybe PrimInfo
forall a b. (a -> b) -> a -> b
$ Var -> TypedArg -> PrimInfo
AdvApp Var
f (Var
argV, Type
argT)
validatePartialPrimInfo PartialPrimInfo
_ = Maybe PrimInfo
forall a. Maybe a
Nothing
isPrimExpr :: Expr Var -> Maybe PrimInfo
isPrimExpr :: Expr Var -> Maybe PrimInfo
isPrimExpr Expr Var
expr = Expr Var -> Maybe PartialPrimInfo
isPrimExpr' Expr Var
expr Maybe PartialPrimInfo
-> (PartialPrimInfo -> Maybe PrimInfo) -> Maybe PrimInfo
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= PartialPrimInfo -> Maybe PrimInfo
validatePartialPrimInfo
isPrimExpr' :: Expr Var -> Maybe PartialPrimInfo
isPrimExpr' :: Expr Var -> Maybe PartialPrimInfo
isPrimExpr' (App Expr Var
e (Type Type
t)) = case Maybe PartialPrimInfo
mPPI of
Just pPI :: PartialPrimInfo
pPI@(PartialPrimInfo {typeArgs :: PartialPrimInfo -> [Type]
typeArgs = [Type]
tArgs}) -> PartialPrimInfo -> Maybe PartialPrimInfo
forall a. a -> Maybe a
Just PartialPrimInfo
pPI {typeArgs = t : tArgs}
Maybe PartialPrimInfo
Nothing -> Maybe PartialPrimInfo
forall a. Maybe a
Nothing
where mPPI :: Maybe PartialPrimInfo
mPPI = Expr Var -> Maybe PartialPrimInfo
isPrimExpr' Expr Var
e
isPrimExpr' (App Expr Var
e Expr Var
e') =
case Expr Var -> Maybe PartialPrimInfo
isPrimExpr' Expr Var
e of
Just partPrimInfo :: PartialPrimInfo
partPrimInfo@(PartialPrimInfo { primPart :: PartialPrimInfo -> Prim
primPart = Prim
Delay, args :: PartialPrimInfo -> [Var]
args = [Var]
args}) -> PartialPrimInfo -> Maybe PartialPrimInfo
forall a. a -> Maybe a
Just PartialPrimInfo
partPrimInfo {args = undefined : args}
Just partPrimInfo :: PartialPrimInfo
partPrimInfo@(PartialPrimInfo { args :: PartialPrimInfo -> [Var]
args = [Var]
args}) -> PartialPrimInfo -> Maybe PartialPrimInfo
forall a. a -> Maybe a
Just PartialPrimInfo
partPrimInfo {args = maybe args (:args) (getMaybeVar e')}
Maybe PartialPrimInfo
_ -> Maybe PartialPrimInfo
forall a. Maybe a
Nothing
isPrimExpr' (Var Var
v) = case Var -> Maybe Prim
isPrim Var
v of
Just Prim
p -> PartialPrimInfo -> Maybe PartialPrimInfo
forall a. a -> Maybe a
Just (PartialPrimInfo -> Maybe PartialPrimInfo)
-> PartialPrimInfo -> Maybe PartialPrimInfo
forall a b. (a -> b) -> a -> b
$ Prim -> Var -> PartialPrimInfo
createPartialPrimInfo Prim
p Var
v
Maybe Prim
Nothing -> Maybe PartialPrimInfo
forall a. Maybe a
Nothing
isPrimExpr' (Tick CoreTickish
_ Expr Var
e) = Expr Var -> Maybe PartialPrimInfo
isPrimExpr' Expr Var
e
isPrimExpr' (Lam Var
v Expr Var
e)
| Var -> Bool
isTyVar Var
v Bool -> Bool -> Bool
|| (Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ Type -> Bool
tcIsLiftedTypeKind (Type -> Bool) -> Type -> Bool
forall a b. (a -> b) -> a -> b
$ (() :: Constraint) => Type -> Type
Type -> Type
typeKind (Type -> Type) -> Type -> Type
forall a b. (a -> b) -> a -> b
$ Var -> Type
varType Var
v) = Expr Var -> Maybe PartialPrimInfo
isPrimExpr' Expr Var
e
isPrimExpr' Expr Var
_ = Maybe PartialPrimInfo
forall a. Maybe a
Nothing