{-# 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

-- DelayApp has the following fields: Var = delay f, T1 = value type, T2 = later type (O v a)
-- AdvApp has the following fields: Var = adv f, TypedArg = var and type for arg
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