{-# 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
text String
"PartialPrimInfo { prim = Delay, function = " SDoc -> SDoc -> SDoc
<> forall a. Outputable a => a -> SDoc
ppr Var
f SDoc -> SDoc -> SDoc
<> String -> SDoc
text String
"args = (not printing since it should be undefined) , typeArgs = " SDoc -> SDoc -> SDoc
<> forall a. Outputable a => a -> SDoc
ppr [Type]
typeArgs
ppr (PartialPrimInfo Prim
prim Var
f [Var]
args [Type]
typeArgs) = String -> SDoc
text String
"PartialPrimInfo { prim = " SDoc -> SDoc -> SDoc
<> forall a. Outputable a => a -> SDoc
ppr Prim
prim SDoc -> SDoc -> SDoc
<> String -> SDoc
text String
", function = " SDoc -> SDoc -> SDoc
<> forall a. Outputable a => a -> SDoc
ppr Var
f SDoc -> SDoc -> SDoc
<> String -> SDoc
text String
", args = " SDoc -> SDoc -> SDoc
<> forall a. Outputable a => a -> SDoc
ppr [Var]
args SDoc -> SDoc -> SDoc
<> String -> SDoc
text String
", typeArgs = " SDoc -> SDoc -> 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
text String
"DelayApp - function " SDoc -> SDoc -> SDoc
<> forall a. Outputable a => a -> SDoc
ppr Var
f
ppr (BoxApp Var
f) = String -> SDoc
text String
"BoxApp - function " SDoc -> SDoc -> SDoc
<> forall a. Outputable a => a -> SDoc
ppr Var
f
ppr (AdvApp Var
f TypedArg
arg) = String -> SDoc
text String
"AdvApp - function " SDoc -> SDoc -> SDoc
<> forall a. Outputable a => a -> SDoc
ppr Var
f SDoc -> SDoc -> SDoc
<> String -> SDoc
text String
" | arg " SDoc -> SDoc -> SDoc
<> forall a. Outputable a => a -> SDoc
ppr TypedArg
arg
ppr (SelectApp Var
f TypedArg
arg TypedArg
arg2) = String -> SDoc
text String
"SelectApp - function " SDoc -> SDoc -> SDoc
<> forall a. Outputable a => a -> SDoc
ppr Var
f SDoc -> SDoc -> SDoc
<> String -> SDoc
text String
" | arg " SDoc -> SDoc -> SDoc
<> forall a. Outputable a => a -> SDoc
ppr TypedArg
arg SDoc -> SDoc -> SDoc
<> String -> SDoc
text String
" | arg2 " SDoc -> SDoc -> SDoc
<> forall a. Outputable a => a -> SDoc
ppr TypedArg
arg2
primMap :: Map FastString Prim
primMap :: Map FastString Prim
primMap = 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 forall a. NamedThing a => a -> Maybe (FastString, FastString)
getNameModule Var
v of
Just (FastString
name, FastString
mod) | FastString -> Bool
isRattModule FastString
mod -> forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup FastString
name Map FastString Prim
primMap
Maybe (FastString, FastString)
_ -> 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]) = forall a. a -> Maybe a
Just 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]) = forall a. a -> Maybe a
Just 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}) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Var -> PrimInfo
BoxApp Var
f
validatePartialPrimInfo (PartialPrimInfo Prim
Adv Var
f [Var
argV] [Type
argT]) = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Var -> TypedArg -> PrimInfo
AdvApp Var
f (Var
argV, Type
argT)
validatePartialPrimInfo PartialPrimInfo
_ = 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 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}) -> forall a. a -> Maybe a
Just PartialPrimInfo
pPI {typeArgs :: [Type]
typeArgs = Type
t forall a. a -> [a] -> [a]
: [Type]
tArgs}
Maybe PartialPrimInfo
Nothing -> 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}) -> forall a. a -> Maybe a
Just PartialPrimInfo
partPrimInfo {args :: [Var]
args = forall a. HasCallStack => a
undefined forall a. a -> [a] -> [a]
: [Var]
args}
Just partPrimInfo :: PartialPrimInfo
partPrimInfo@(PartialPrimInfo { args :: PartialPrimInfo -> [Var]
args = [Var]
args}) -> forall a. a -> Maybe a
Just PartialPrimInfo
partPrimInfo {args :: [Var]
args = forall b a. b -> (a -> b) -> Maybe a -> b
maybe [Var]
args (forall a. a -> [a] -> [a]
:[Var]
args) (Expr Var -> Maybe Var
getMaybeVar Expr Var
e')}
Maybe PartialPrimInfo
_ -> forall a. Maybe a
Nothing
isPrimExpr' (Var Var
v) = case Var -> Maybe Prim
isPrim Var
v of
Just Prim
p -> forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Prim -> Var -> PartialPrimInfo
createPartialPrimInfo Prim
p Var
v
Maybe Prim
Nothing -> 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 forall a b. (a -> b) -> a -> b
$ Type -> Bool
tcIsLiftedTypeKind forall a b. (a -> b) -> a -> b
$ HasDebugCallStack => Type -> Type
typeKind forall a b. (a -> b) -> a -> b
$ Var -> Type
varType Var
v) = Expr Var -> Maybe PartialPrimInfo
isPrimExpr' Expr Var
e
isPrimExpr' Expr Var
_ = forall a. Maybe a
Nothing