module Test.Inspection.Core
( slice
, pprSlice
, pprSliceDifference
, eqSlice
, freeOfType
, doesNotAllocate
) where
import CoreSyn
import CoreUtils
import TyCoRep
import Type
import Var
import Id
import Name
import VarEnv
import Literal (nullAddrLit)
import Outputable
import PprCore
import Coercion
import Util
import qualified Data.Set as S
import Data.Maybe
import State
import Control.Monad
type Slice = [(Var, CoreExpr)]
slice :: [(Var, CoreExpr)] -> Var -> Slice
slice binds v = [(v,e) | (v,e) <- binds, v `S.member` used ]
where
used = execState (goV v) S.empty
local = S.fromList (map fst binds)
goV v | v `S.member` local = do
seen <- gets (v `S.member`)
unless seen $ do
modify (S.insert v)
let Just e = lookup v binds
go e
| otherwise = return ()
go (Var v) = goV v
go (Lit _ ) = pure ()
go (App e arg) | isTypeArg arg = go e
go (App e arg) = go e >> go arg
go (Lam b e) | isTyVar b = go e
go (Lam _ e) = go e
go (Let bind body) = mapM_ go (rhssOfBind bind) >> go body
go (Case s _ _ alts) = go s >> mapM_ goA alts
go (Cast e _) = go e
go (Tick _ e) = go e
go (Type _) = pure ()
go (Coercion _) = pure ()
goA (_, _, e) = go e
pprSlice :: Slice -> SDoc
pprSlice slice = withLessDetail $ pprCoreBindings [ NonRec v e | (v,e) <- slice ]
pprSliceDifference :: Slice -> Slice -> SDoc
pprSliceDifference slice1 slice2 =
nest 4 (hang (text "LHS" <> colon) 4 (pprSlice slice1')) $$
nest 4 (hang (text "RHS" <> colon) 4 (pprSlice slice2'))
where
both = S.intersection (S.fromList (map fst slice1)) (S.fromList (map fst slice2))
slice1' = filter (\(v,_) -> v `S.notMember` both) slice1
slice2' = filter (\(v,_) -> v `S.notMember` both) slice2
withLessDetail :: SDoc -> SDoc
#if MIN_VERSION_GLASGOW_HASKELL(8,2,0,0)
withLessDetail sdoc = sdocWithDynFlags $ \dflags ->
withPprStyle (defaultUserStyle dflags) sdoc
#else
withLessDetail sdoc = withPprStyle defaultUserStyle sdoc
#endif
eqSlice :: Bool -> Slice -> Slice -> Bool
eqSlice it slice1 slice2
= go (mkRnEnv2 emptyInScopeSet)
(Let (Rec slice1) (Lit nullAddrLit))
(Let (Rec slice2) (Lit nullAddrLit))
where
go env (Var v1) (Var v2)
| rnOccL env v1 == rnOccR env v2
= True
go _ (Lit lit1) (Lit lit2) = lit1 == lit2
go env (Type t1) (Type t2) = eqTypeX env t1 t2
go env (Coercion co1) (Coercion co2) = eqCoercionX env co1 co2
go env (Cast e1 _) e2 | it = go env e1 e2
go env e1 (Cast e2 _) | it = go env e1 e2
go env (Cast e1 co1) (Cast e2 co2) = eqCoercionX env co1 co2 && go env e1 e2
go env (App e1 a) e2 | it, isTypeArg a = go env e1 e2
go env e1 (App e2 a) | it, isTypeArg a = go env e1 e2
go env (App f1 a1) (App f2 a2) = go env f1 f2 && go env a1 a2
go env (Tick n1 e1) (Tick n2 e2) = go_tick env n1 n2 && go env e1 e2
go env (Lam b e1) e2 | it, isTyCoVar b = go env e1 e2
go env e1 (Lam b e2) | it, isTyCoVar b = go env e1 e2
go env (Lam b1 e1) (Lam b2 e2)
= (it || eqTypeX env (varType b1) (varType b2))
&& go (rnBndr2 env b1 b2) e1 e2
go env (Let (NonRec v1 r1) e1) (Let (NonRec v2 r2) e2)
= go env r1 r2
&& go (rnBndr2 env v1 v2) e1 e2
go env (Let (Rec ps1) e1) (Let (Rec ps2) e2)
= equalLength ps1 ps2
&& all2 (go env') rs1 rs2 && go env' e1 e2
where
(bs1,rs1) = unzip ps1
(bs2,rs2) = unzip ps2
env' = rnBndrs2 env bs1 bs2
go env (Case e1 b1 t1 a1) (Case e2 b2 t2 a2)
| null a1
= null a2 && go env e1 e2 && (it || eqTypeX env t1 t2)
| otherwise
= go env e1 e2 && all2 (go_alt (rnBndr2 env b1 b2)) a1 a2
go _ _ _ = False
go_alt env (c1, bs1, e1) (c2, bs2, e2)
= c1 == c2 && go (rnBndrs2 env bs1 bs2) e1 e2
go_tick :: RnEnv2 -> Tickish Id -> Tickish Id -> Bool
go_tick env (Breakpoint lid lids) (Breakpoint rid rids)
= lid == rid && map (rnOccL env) lids == map (rnOccR env) rids
go_tick _ l r = l == r
freeOfType :: Slice -> Name -> Maybe (Var, CoreExpr)
freeOfType slice tcN = listToMaybe [ (v,e) | (v,e) <- slice, not (go e) ]
where
goV v = goT (varType v)
go (Var v) = goV v
go (Lit _ ) = True
go (App e a) = go e && go a
go (Lam b e) = goV b && go e
go (Let bind body) = all goB (flattenBinds [bind]) && go body
go (Case s b _ alts) = go s && goV b && all goA alts
go (Cast e _) = go e
go (Tick _ e) = go e
go (Type t) = (goT t)
go (Coercion _) = True
goB (b, e) = goV b && go e
goA (_,pats, e) = all goV pats && go e
goT (TyVarTy _) = True
goT (AppTy t1 t2) = goT t1 && goT t2
goT (TyConApp tc ts) = getName tc /= tcN && all goT ts
goT (ForAllTy _ t) = goT t
#if MIN_VERSION_GLASGOW_HASKELL(8,2,0,0)
goT (FunTy t1 t2) = goT t1 && goT t2
#endif
goT (LitTy _) = True
goT (CastTy t _) = goT t
goT (CoercionTy _) = True
doesNotAllocate :: Slice -> Maybe (Var, CoreExpr)
doesNotAllocate slice = listToMaybe [ (v,e) | (v,e) <- slice, not (go (idArity v) e) ]
where
go _ (Var v)
| isDataConWorkId v, idArity v > 0 = False
go a (Var v) = (a >= idArity v)
go _ (Lit _ ) = True
go a (App e arg) | isTypeArg arg = go a e
go a (App e arg) = go (a+1) e && goArg arg
go a (Lam b e) | isTyVar b = go a e
go 0 (Lam _ _) = False
go a (Lam _ e) = go (a1) e
go a (Let bind body) = all goB (flattenBinds [bind]) && go a body
go a (Case s _ _ alts) = go 0 s && all (goA a) alts
go a (Cast e _) = go a e
go a (Tick _ e) = go a e
go _ (Type _) = True
go _ (Coercion _) = True
goArg e | exprIsTrivial e = go 0 e
| isUnliftedType (exprType e) = go 0 e
| otherwise = False
goB (b, e)
#if MIN_VERSION_GLASGOW_HASKELL(8,2,0,0)
| isJoinId b = go (idArity b) e
#endif
| isFunTy (idType b) = go (idArity b) e
| isUnliftedType (idType b) = go (idArity b) e
| otherwise = False
goA a (_,_, e) = go a e