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