{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
module Rattus.Plugin.Utils (
  printMessage,
  Severity(..),
  isRattModule,
  isGhcModule,
  getNameModule,
  isStable,
  isStrict,
  isTemporal,
  userFunction,
  isType)
  where

import ErrUtils
import Prelude hiding ((<>))
import GhcPlugins
import Data.Set (Set)
import qualified Data.Set as Set
import Data.Char
import Data.Maybe

isType :: Expr b -> Bool
isType Type {} = Bool
True
isType (App e :: Expr b
e _) = Expr b -> Bool
isType Expr b
e
isType (Cast e :: Expr b
e _) = Expr b -> Bool
isType Expr b
e
isType (Tick _ e :: Expr b
e) = Expr b -> Bool
isType Expr b
e
isType _ = Bool
False



-- printMessage :: Severity -> SDoc -> CoreM ()
-- printMessage sev doc =
--   case sev of
--     Error -> errorMsg doc
--     Warning -> warnMsg doc


-- printMessageV :: Severity -> Var -> SDoc -> CoreM ()
-- printMessageV sev var doc =
--   let loc = nameSrcLoc (varName var)
--       doc' = ppr loc <> text ": " <> doc
--   in case sev of
--     Error -> errorMsg doc'
--     Warning -> warnMsg doc'


printMessage :: Severity -> SrcSpan -> SDoc -> CoreM ()
printMessage :: Severity -> SrcSpan -> SDoc -> CoreM ()
printMessage sev :: Severity
sev loc :: SrcSpan
loc doc :: SDoc
doc = do
  DynFlags
dflags <- CoreM DynFlags
forall (m :: * -> *). HasDynFlags m => m DynFlags
getDynFlags
  PrintUnqualified
unqual <- CoreM PrintUnqualified
getPrintUnqualified
  let sty :: PprStyle
sty = case Severity
sev of
                     SevError   -> PprStyle
err_sty
                     SevWarning -> PprStyle
err_sty
                     SevDump    -> PprStyle
dump_sty
                     _          -> PprStyle
user_sty
      err_sty :: PprStyle
err_sty  = DynFlags -> PrintUnqualified -> PprStyle
mkErrStyle DynFlags
dflags PrintUnqualified
unqual
      user_sty :: PprStyle
user_sty = DynFlags -> PrintUnqualified -> Depth -> PprStyle
mkUserStyle DynFlags
dflags PrintUnqualified
unqual Depth
AllTheWay
      dump_sty :: PprStyle
dump_sty = DynFlags -> PrintUnqualified -> PprStyle
mkDumpStyle DynFlags
dflags PrintUnqualified
unqual
  IO () -> CoreM ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> CoreM ()) -> IO () -> CoreM ()
forall a b. (a -> b) -> a -> b
$ DynFlags
-> WarnReason -> Severity -> SrcSpan -> PprStyle -> SDoc -> IO ()
putLogMsg DynFlags
dflags WarnReason
NoReason Severity
sev SrcSpan
loc PprStyle
sty SDoc
doc



rattModules :: Set FastString
rattModules :: Set FastString
rattModules = [FastString] -> Set FastString
forall a. Ord a => [a] -> Set a
Set.fromList ["Rattus.Internal","Rattus.Primitives"
                           ,"Rattus.Stable", "Rattus.Arrow"]

isRattModule :: FastString -> Bool
isRattModule :: FastString -> Bool
isRattModule = (FastString -> Set FastString -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set FastString
rattModules)

isGhcModule :: FastString -> Bool
isGhcModule :: FastString -> Bool
isGhcModule = (FastString -> FastString -> Bool
forall a. Eq a => a -> a -> Bool
== "GHC.Types")


getNameModule :: NamedThing a => a -> Maybe (FastString, FastString)
getNameModule :: a -> Maybe (FastString, FastString)
getNameModule v :: a
v = do
  let name :: Name
name = a -> Name
forall a. NamedThing a => a -> Name
getName a
v
  Module
mod <- Name -> Maybe Module
nameModule_maybe Name
name
  (FastString, FastString) -> Maybe (FastString, FastString)
forall (m :: * -> *) a. Monad m => a -> m a
return (Name -> FastString
forall a. NamedThing a => a -> FastString
getOccFS Name
name,ModuleName -> FastString
moduleNameFS (Module -> ModuleName
moduleName Module
mod))


-- | The set of stable built-in types.
ghcStableTypes :: Set FastString
ghcStableTypes :: Set FastString
ghcStableTypes = [FastString] -> Set FastString
forall a. Ord a => [a] -> Set a
Set.fromList ["Int","Bool","Float","Double","Char", "IO"]


newtype TypeCmp = TC Type

instance Eq TypeCmp where
  (TC t1 :: Type
t1) == :: TypeCmp -> TypeCmp -> Bool
== (TC t2 :: Type
t2) = Type -> Type -> Bool
eqType Type
t1 Type
t2

instance Ord TypeCmp where
  compare :: TypeCmp -> TypeCmp -> Ordering
compare (TC t1 :: Type
t1) (TC t2 :: Type
t2) = Type -> Type -> Ordering
nonDetCmpType Type
t1 Type
t2

isTemporal :: Type -> Bool
isTemporal :: Type -> Bool
isTemporal t :: Type
t = Int -> Set TypeCmp -> Type -> Bool
isTemporalRec 0 Set TypeCmp
forall a. Set a
Set.empty Type
t


isTemporalRec :: Int -> Set TypeCmp -> Type -> Bool
isTemporalRec :: Int -> Set TypeCmp -> Type -> Bool
isTemporalRec d :: Int
d _ _ | Int
d Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== 100 = Bool
False
isTemporalRec _ pr :: Set TypeCmp
pr t :: Type
t | TypeCmp -> Set TypeCmp -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member (Type -> TypeCmp
TC Type
t) Set TypeCmp
pr = Bool
False
isTemporalRec d :: Int
d pr :: Set TypeCmp
pr t :: Type
t = do
  let pr' :: Set TypeCmp
pr' = TypeCmp -> Set TypeCmp -> Set TypeCmp
forall a. Ord a => a -> Set a -> Set a
Set.insert (Type -> TypeCmp
TC Type
t) Set TypeCmp
pr
  case HasDebugCallStack => Type -> Maybe (TyCon, [Type])
Type -> Maybe (TyCon, [Type])
splitTyConApp_maybe Type
t of
    Nothing -> Bool
False
    Just (con :: TyCon
con,args :: [Type]
args) ->
      case TyCon -> Maybe (FastString, FastString)
forall a. NamedThing a => a -> Maybe (FastString, FastString)
getNameModule TyCon
con of
        Nothing -> Bool
False
        Just (name :: FastString
name,mod :: FastString
mod)
          -- If it's a Rattus type constructor check if it's a box
          | FastString -> Bool
isRattModule FastString
mod Bool -> Bool -> Bool
&& (FastString
name FastString -> FastString -> Bool
forall a. Eq a => a -> a -> Bool
== "Box" Bool -> Bool -> Bool
|| FastString
name FastString -> FastString -> Bool
forall a. Eq a => a -> a -> Bool
== "O") -> Bool
True
          | TyCon -> Bool
isFunTyCon TyCon
con -> [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
or ((Type -> Bool) -> [Type] -> [Bool]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Set TypeCmp -> Type -> Bool
isTemporalRec (Int
dInt -> Int -> Int
forall a. Num a => a -> a -> a
+1) Set TypeCmp
pr') [Type]
args)
          | TyCon -> Bool
isAlgTyCon TyCon
con -> 
            case TyCon -> AlgTyConRhs
algTyConRhs TyCon
con of
              DataTyCon {data_cons :: AlgTyConRhs -> [DataCon]
data_cons = [DataCon]
cons} -> [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
or ((DataCon -> Bool) -> [DataCon] -> [Bool]
forall a b. (a -> b) -> [a] -> [b]
map DataCon -> Bool
check [DataCon]
cons)
                where check :: DataCon -> Bool
check con :: DataCon
con = case DataCon -> [Type] -> ([TyCoVar], [Type], [Type])
dataConInstSig DataCon
con [Type]
args of
                        (_, _,tys :: [Type]
tys) -> [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
or ((Type -> Bool) -> [Type] -> [Bool]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Set TypeCmp -> Type -> Bool
isTemporalRec (Int
dInt -> Int -> Int
forall a. Num a => a -> a -> a
+1) Set TypeCmp
pr') [Type]
tys)
              _ -> [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
or ((Type -> Bool) -> [Type] -> [Bool]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Set TypeCmp -> Type -> Bool
isTemporalRec (Int
dInt -> Int -> Int
forall a. Num a => a -> a -> a
+1) Set TypeCmp
pr') [Type]
args)
        _ -> Bool
False


-- | Check whether the given type is stable. This check may use
-- 'Stable' constraints from the context.

isStable :: Set Var -> Type -> Bool
isStable :: Set TyCoVar -> Type -> Bool
isStable c :: Set TyCoVar
c t :: Type
t = Set TyCoVar -> Int -> Set TypeCmp -> Type -> Bool
isStableRec Set TyCoVar
c 0 Set TypeCmp
forall a. Set a
Set.empty Type
t

-- | Check whether the given type is stable. This check may use
-- 'Stable' constraints from the context.

isStableRec :: Set Var -> Int -> Set TypeCmp -> Type -> Bool
-- To prevent infinite recursion (when checking recursive types) we
-- keep track of previously checked types. This, however, is not
-- enough for non-regular data types. Hence we also have a counter.
isStableRec :: Set TyCoVar -> Int -> Set TypeCmp -> Type -> Bool
isStableRec _ d :: Int
d _ _ | Int
d Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== 100 = Bool
True
isStableRec _ _ pr :: Set TypeCmp
pr t :: Type
t | TypeCmp -> Set TypeCmp -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member (Type -> TypeCmp
TC Type
t) Set TypeCmp
pr = Bool
True
isStableRec c :: Set TyCoVar
c d :: Int
d pr :: Set TypeCmp
pr t :: Type
t = do
  let pr' :: Set TypeCmp
pr' = TypeCmp -> Set TypeCmp -> Set TypeCmp
forall a. Ord a => a -> Set a -> Set a
Set.insert (Type -> TypeCmp
TC Type
t) Set TypeCmp
pr
  case HasDebugCallStack => Type -> Maybe (TyCon, [Type])
Type -> Maybe (TyCon, [Type])
splitTyConApp_maybe Type
t of
    Nothing -> case Type -> Maybe TyCoVar
getTyVar_maybe Type
t of
      Just v :: TyCoVar
v -> -- if it's a type variable, check the context
        TyCoVar
v TyCoVar -> Set TyCoVar -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set TyCoVar
c
      Nothing -> Bool
False
    Just (con :: TyCon
con,args :: [Type]
args) ->
      case TyCon -> Maybe (FastString, FastString)
forall a. NamedThing a => a -> Maybe (FastString, FastString)
getNameModule TyCon
con of
        Nothing -> Bool
False
        Just (name :: FastString
name,mod :: FastString
mod)
          -- If it's a Rattus type constructor check if it's a box
          | FastString -> Bool
isRattModule FastString
mod Bool -> Bool -> Bool
&& FastString
name FastString -> FastString -> Bool
forall a. Eq a => a -> a -> Bool
== "Box" -> Bool
True
            -- If its a built-in type check the set of stable built-in types
          | FastString -> Bool
isGhcModule FastString
mod -> FastString
name FastString -> Set FastString -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set FastString
ghcStableTypes
          {- deal with type synonyms (does not seem to be necessary (??))
           | Just (subst,ty,[]) <- expandSynTyCon_maybe con args ->
             isStableRec c (d+1) pr' (substTy (extendTvSubstList emptySubst subst) ty) -}
          | TyCon -> Bool
isAlgTyCon TyCon
con -> 
            case TyCon -> AlgTyConRhs
algTyConRhs TyCon
con of
              DataTyCon {data_cons :: AlgTyConRhs -> [DataCon]
data_cons = [DataCon]
cons, is_enum :: AlgTyConRhs -> Bool
is_enum = Bool
enum}
                | Bool
enum -> Bool
True
                | [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ([Bool] -> Bool) -> [Bool] -> Bool
forall a b. (a -> b) -> a -> b
$ (DataCon -> [Bool]) -> [DataCon] -> [Bool]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ((HsSrcBang -> Bool) -> [HsSrcBang] -> [Bool]
forall a b. (a -> b) -> [a] -> [b]
map HsSrcBang -> Bool
isSrcStrict'
                                   ([HsSrcBang] -> [Bool])
-> (DataCon -> [HsSrcBang]) -> DataCon -> [Bool]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DataCon -> [HsSrcBang]
dataConSrcBangs) ([DataCon] -> [Bool]) -> [DataCon] -> [Bool]
forall a b. (a -> b) -> a -> b
$ [DataCon]
cons ->
                  [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and  ((DataCon -> Bool) -> [DataCon] -> [Bool]
forall a b. (a -> b) -> [a] -> [b]
map DataCon -> Bool
check [DataCon]
cons)
                | Bool
otherwise -> Bool
False
                where check :: DataCon -> Bool
check con :: DataCon
con = case DataCon -> [Type] -> ([TyCoVar], [Type], [Type])
dataConInstSig DataCon
con [Type]
args of
                        (_, _,tys :: [Type]
tys) -> [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ((Type -> Bool) -> [Type] -> [Bool]
forall a b. (a -> b) -> [a] -> [b]
map (Set TyCoVar -> Int -> Set TypeCmp -> Type -> Bool
isStableRec Set TyCoVar
c (Int
dInt -> Int -> Int
forall a. Num a => a -> a -> a
+1) Set TypeCmp
pr') [Type]
tys)
              TupleTyCon {} -> [Type] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Type]
args
              _ -> Bool
False
        _ -> Bool
False



isStrict :: Type -> Bool
isStrict :: Type -> Bool
isStrict t :: Type
t = Int -> Set TypeCmp -> Type -> Bool
isStrictRec 0 Set TypeCmp
forall a. Set a
Set.empty Type
t

-- | Check whether the given type is stable. This check may use
-- 'Stable' constraints from the context.

isStrictRec :: Int -> Set TypeCmp -> Type -> Bool
-- To prevent infinite recursion (when checking recursive types) we
-- keep track of previously checked types. This, however, is not
-- enough for non-regular data types. Hence we also have a counter.
isStrictRec :: Int -> Set TypeCmp -> Type -> Bool
isStrictRec d :: Int
d _ _ | Int
d Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== 100 = Bool
True
isStrictRec _ pr :: Set TypeCmp
pr t :: Type
t | TypeCmp -> Set TypeCmp -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member (Type -> TypeCmp
TC Type
t) Set TypeCmp
pr = Bool
True
isStrictRec d :: Int
d pr :: Set TypeCmp
pr t :: Type
t = do
  let pr' :: Set TypeCmp
pr' = TypeCmp -> Set TypeCmp -> Set TypeCmp
forall a. Ord a => a -> Set a -> Set a
Set.insert (Type -> TypeCmp
TC Type
t) Set TypeCmp
pr
  let (_,t' :: Type
t') = Type -> ([TyCoVar], Type)
splitForAllTys Type
t
  let (c :: Type
c, tys :: [Type]
tys) = HasDebugCallStack => Type -> (Type, [Type])
Type -> (Type, [Type])
repSplitAppTys Type
t'
  if Maybe TyCoVar -> Bool
forall a. Maybe a -> Bool
isJust (Type -> Maybe TyCoVar
getTyVar_maybe Type
c) then [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ((Type -> Bool) -> [Type] -> [Bool]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Set TypeCmp -> Type -> Bool
isStrictRec (Int
dInt -> Int -> Int
forall a. Num a => a -> a -> a
+1) Set TypeCmp
pr') [Type]
tys)
  else  case HasDebugCallStack => Type -> Maybe (TyCon, [Type])
Type -> Maybe (TyCon, [Type])
splitTyConApp_maybe Type
t' of
    Nothing -> Maybe TyCoVar -> Bool
forall a. Maybe a -> Bool
isJust (Type -> Maybe TyCoVar
getTyVar_maybe Type
t)
    Just (con :: TyCon
con,args :: [Type]
args) ->
      case TyCon -> Maybe (FastString, FastString)
forall a. NamedThing a => a -> Maybe (FastString, FastString)
getNameModule TyCon
con of
        Nothing -> Bool
False
        Just (name :: FastString
name,mod :: FastString
mod)
          -- If it's a Rattus type constructor check if it's a box
          | FastString -> Bool
isRattModule FastString
mod Bool -> Bool -> Bool
&& (FastString
name FastString -> FastString -> Bool
forall a. Eq a => a -> a -> Bool
== "Box" Bool -> Bool -> Bool
|| FastString
name FastString -> FastString -> Bool
forall a. Eq a => a -> a -> Bool
== "O") -> Bool
True
            -- If its a built-in type check the set of stable built-in types
          | FastString -> Bool
isGhcModule FastString
mod -> FastString
name FastString -> Set FastString -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set FastString
ghcStableTypes
          {- deal with type synonyms (does not seem to be necessary (??))
           | Just (subst,ty,[]) <- expandSynTyCon_maybe con args ->
             isStrictRec c (d+1) pr' (substTy (extendTvSubstList emptySubst subst) ty) -}
          | TyCon -> Bool
isFunTyCon TyCon
con -> Bool
True
          | TyCon -> Bool
isAlgTyCon TyCon
con -> 
            case TyCon -> AlgTyConRhs
algTyConRhs TyCon
con of
              DataTyCon {data_cons :: AlgTyConRhs -> [DataCon]
data_cons = [DataCon]
cons, is_enum :: AlgTyConRhs -> Bool
is_enum = Bool
enum}
                | Bool
enum -> Bool
True
                | [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ([Bool] -> Bool) -> [Bool] -> Bool
forall a b. (a -> b) -> a -> b
$ ((DataCon -> Bool) -> [DataCon] -> [Bool]
forall a b. (a -> b) -> [a] -> [b]
map ([Type] -> DataCon -> Bool
isSrcStrictOrDelay [Type]
args)) ([DataCon] -> [Bool]) -> [DataCon] -> [Bool]
forall a b. (a -> b) -> a -> b
$ [DataCon]
cons ->
                  [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and  ((DataCon -> Bool) -> [DataCon] -> [Bool]
forall a b. (a -> b) -> [a] -> [b]
map DataCon -> Bool
check [DataCon]
cons)
                | Bool
otherwise -> Bool
False
                where check :: DataCon -> Bool
check con :: DataCon
con = case DataCon -> [Type] -> ([TyCoVar], [Type], [Type])
dataConInstSig DataCon
con [Type]
args of
                        (_, _,tys :: [Type]
tys) -> [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ((Type -> Bool) -> [Type] -> [Bool]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Set TypeCmp -> Type -> Bool
isStrictRec (Int
dInt -> Int -> Int
forall a. Num a => a -> a -> a
+1) Set TypeCmp
pr') [Type]
tys)
              TupleTyCon {} -> [Type] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Type]
args
              _ -> Bool
False
          | Bool
otherwise -> Bool
False
            




isSrcStrictOrDelay :: [Type] -> DataCon -> Bool
isSrcStrictOrDelay :: [Type] -> DataCon -> Bool
isSrcStrictOrDelay args :: [Type]
args con :: DataCon
con = [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and ((Type -> HsSrcBang -> Bool) -> [Type] -> [HsSrcBang] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Type -> HsSrcBang -> Bool
check [Type]
tys (DataCon -> [HsSrcBang]
dataConSrcBangs DataCon
con))
  where (_, _,tys :: [Type]
tys) = DataCon -> [Type] -> ([TyCoVar], [Type], [Type])
dataConInstSig DataCon
con [Type]
args 
        check :: Type -> HsSrcBang -> Bool
check ty :: Type
ty b :: HsSrcBang
b = HsSrcBang -> Bool
isSrcStrict' HsSrcBang
b Bool -> Bool -> Bool
|| Type -> Bool
isDelay Type
ty
        isDelay :: Type -> Bool
isDelay ty :: Type
ty = case HasDebugCallStack => Type -> Maybe (TyCon, [Type])
Type -> Maybe (TyCon, [Type])
splitTyConApp_maybe Type
ty of
                       Just (con :: TyCon
con,_) ->
                         case TyCon -> Maybe (FastString, FastString)
forall a. NamedThing a => a -> Maybe (FastString, FastString)
getNameModule TyCon
con of
                           Just (name :: FastString
name,mod :: FastString
mod) | FastString -> Bool
isRattModule FastString
mod Bool -> Bool -> Bool
&& FastString
name FastString -> FastString -> Bool
forall a. Eq a => a -> a -> Bool
== "O" -> Bool
True
                           _ -> Bool
False
                       _ -> Bool
False

isSrcStrict' :: HsSrcBang -> Bool
isSrcStrict' (HsSrcBang _ _ SrcStrict) = Bool
True
isSrcStrict' _ = Bool
False


userFunction :: Var -> Bool
userFunction :: TyCoVar -> Bool
userFunction v :: TyCoVar
v =
  case Name -> String
forall a. NamedThing a => a -> String
getOccString (TyCoVar -> Name
forall a. NamedThing a => a -> Name
getName TyCoVar
v) of
    (c :: Char
c : _)
      | Char -> Bool
isUpper Char
c Bool -> Bool -> Bool
|| Char
c Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== '$' Bool -> Bool -> Bool
|| Char
c Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== ':' -> Bool
False
      | Bool
otherwise -> Bool
True
    _ -> Bool
False