{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE FlexibleInstances         #-}
{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE ScopedTypeVariables       #-}

module Language.Haskell.Liquid.Types.Visitors (

  CBVisitable (..)

  -- * visitors
  , coreVisitor
  , CoreVisitor (..)

  ) where

import           Data.Hashable

import           Data.List                        (foldl', (\\), delete)
import qualified Data.HashSet                     as S
import           Prelude                          hiding (error)
import           Language.Fixpoint.Misc
import           Liquid.GHC.API
import           Language.Haskell.Liquid.GHC.Misc ()


------------------------------------------------------------------------------
-------------------------------- A CoreBind Visitor --------------------------
------------------------------------------------------------------------------

-- TODO: syb-shrinkage

class CBVisitable a where
  freeVars :: S.HashSet Var -> a -> [Var]
  readVars :: a -> [Var]
  letVars  :: a -> [Var]
  literals :: a -> [Literal]

instance CBVisitable [CoreBind] where
  freeVars :: HashSet Var -> [CoreBind] -> [Var]
freeVars HashSet Var
env [CoreBind]
cbs = forall a. Ord a => [a] -> [a]
sortNub [Var]
xs forall a. Eq a => [a] -> [a] -> [a]
\\ [Var]
ys
    where xs :: [Var]
xs = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (forall a. CBVisitable a => HashSet Var -> a -> [Var]
freeVars HashSet Var
env) [CoreBind]
cbs
          ys :: [Var]
ys = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall t. Bind t -> [t]
bindings [CoreBind]
cbs

  readVars :: [CoreBind] -> [Var]
readVars = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall a. CBVisitable a => a -> [Var]
readVars
  letVars :: [CoreBind] -> [Var]
letVars  = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall a. CBVisitable a => a -> [Var]
letVars
  literals :: [CoreBind] -> [Literal]
literals = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall a. CBVisitable a => a -> [Literal]
literals

instance CBVisitable CoreBind where
  freeVars :: HashSet Var -> CoreBind -> [Var]
freeVars HashSet Var
env (NonRec Var
x Expr Var
e) = forall a. CBVisitable a => HashSet Var -> a -> [Var]
freeVars (forall a. (Eq a, Hashable a) => HashSet a -> [a] -> HashSet a
extendEnv HashSet Var
env [Var
x]) Expr Var
e
  freeVars HashSet Var
env (Rec [(Var, Expr Var)]
xes)    = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (forall a. CBVisitable a => HashSet Var -> a -> [Var]
freeVars HashSet Var
env') [Expr Var]
es
                              where ([Var]
xs,[Expr Var]
es) = forall a b. [(a, b)] -> ([a], [b])
unzip [(Var, Expr Var)]
xes
                                    env' :: HashSet Var
env'    = forall a. (Eq a, Hashable a) => HashSet a -> [a] -> HashSet a
extendEnv HashSet Var
env [Var]
xs

  readVars :: CoreBind -> [Var]
readVars (NonRec Var
_ Expr Var
e)     = forall a. CBVisitable a => a -> [Var]
readVars Expr Var
e
  readVars (Rec [(Var, Expr Var)]
xes)        = forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [Var
x forall a. Eq a => a -> [a] -> [a]
`delete` forall a. CBVisitable a => a -> [Var]
nubReadVars Expr Var
e |(Var
x, Expr Var
e) <- [(Var, Expr Var)]
xes]
    where nubReadVars :: a -> [Var]
nubReadVars = forall a. Ord a => [a] -> [a]
sortNub forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. CBVisitable a => a -> [Var]
readVars

  letVars :: CoreBind -> [Var]
letVars (NonRec Var
x Expr Var
e)      = Var
x forall a. a -> [a] -> [a]
: forall a. CBVisitable a => a -> [Var]
letVars Expr Var
e
  letVars (Rec [(Var, Expr Var)]
xes)         = [Var]
xs forall a. [a] -> [a] -> [a]
++ forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall a. CBVisitable a => a -> [Var]
letVars [Expr Var]
es
    where
      ([Var]
xs, [Expr Var]
es)              = forall a b. [(a, b)] -> ([a], [b])
unzip [(Var, Expr Var)]
xes

  literals :: CoreBind -> [Literal]
literals (NonRec Var
_ Expr Var
e)      = forall a. CBVisitable a => a -> [Literal]
literals Expr Var
e
  literals (Rec [(Var, Expr Var)]
xes)         = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (forall a. CBVisitable a => a -> [Literal]
literals forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) [(Var, Expr Var)]
xes

instance CBVisitable (Expr Var) where
  freeVars :: HashSet Var -> Expr Var -> [Var]
freeVars = HashSet Var -> Expr Var -> [Var]
exprFreeVars
  readVars :: Expr Var -> [Var]
readVars = forall t.
(CBVisitable (Alt t), CBVisitable (Bind t)) =>
Expr t -> [Var]
exprReadVars
  letVars :: Expr Var -> [Var]
letVars  = Expr Var -> [Var]
exprLetVars
  literals :: Expr Var -> [Literal]
literals = forall t.
(CBVisitable (Alt t), CBVisitable (Bind t)) =>
Expr t -> [Literal]
exprLiterals

exprFreeVars :: S.HashSet Id -> Expr Id -> [Id]
exprFreeVars :: HashSet Var -> Expr Var -> [Var]
exprFreeVars = HashSet Var -> Expr Var -> [Var]
go
  where
    go :: HashSet Var -> Expr Var -> [Var]
go HashSet Var
env (Var Var
x)         = [Var
x | Bool -> Bool
not (Var
x forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
`S.member` HashSet Var
env)]
    go HashSet Var
env (App Expr Var
e Expr Var
a)       = HashSet Var -> Expr Var -> [Var]
go HashSet Var
env Expr Var
e forall a. [a] -> [a] -> [a]
++ HashSet Var -> Expr Var -> [Var]
go HashSet Var
env Expr Var
a
    go HashSet Var
env (Lam Var
x Expr Var
e)       = HashSet Var -> Expr Var -> [Var]
go (forall a. (Eq a, Hashable a) => HashSet a -> [a] -> HashSet a
extendEnv HashSet Var
env [Var
x]) Expr Var
e
    go HashSet Var
env (Let CoreBind
b Expr Var
e)       = forall a. CBVisitable a => HashSet Var -> a -> [Var]
freeVars HashSet Var
env CoreBind
b forall a. [a] -> [a] -> [a]
++ HashSet Var -> Expr Var -> [Var]
go (forall a. (Eq a, Hashable a) => HashSet a -> [a] -> HashSet a
extendEnv HashSet Var
env (forall t. Bind t -> [t]
bindings CoreBind
b)) Expr Var
e
    go HashSet Var
env (Tick CoreTickish
_ Expr Var
e)      = HashSet Var -> Expr Var -> [Var]
go HashSet Var
env Expr Var
e
    go HashSet Var
env (Cast Expr Var
e CoercionR
_)      = HashSet Var -> Expr Var -> [Var]
go HashSet Var
env Expr Var
e
    go HashSet Var
env (Case Expr Var
e Var
x Type
_ [Alt Var]
cs) = HashSet Var -> Expr Var -> [Var]
go HashSet Var
env Expr Var
e forall a. [a] -> [a] -> [a]
++ forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (forall a. CBVisitable a => HashSet Var -> a -> [Var]
freeVars (forall a. (Eq a, Hashable a) => HashSet a -> [a] -> HashSet a
extendEnv HashSet Var
env [Var
x])) [Alt Var]
cs
    go HashSet Var
_   Expr Var
_               = []

exprReadVars :: (CBVisitable (Alt t), CBVisitable (Bind t)) => Expr t -> [Id]
exprReadVars :: forall t.
(CBVisitable (Alt t), CBVisitable (Bind t)) =>
Expr t -> [Var]
exprReadVars = forall {b}.
(CBVisitable (Bind b), CBVisitable (Alt b)) =>
Expr b -> [Var]
go
  where
    go :: Expr b -> [Var]
go (Var Var
x)             = [Var
x]
    go (App Expr b
e Expr b
a)           = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Expr b -> [Var]
go [Expr b
e, Expr b
a]
    go (Lam b
_ Expr b
e)           = Expr b -> [Var]
go Expr b
e
    go (Let Bind b
b Expr b
e)           = forall a. CBVisitable a => a -> [Var]
readVars Bind b
b forall a. [a] -> [a] -> [a]
++ Expr b -> [Var]
go Expr b
e
    go (Tick CoreTickish
_ Expr b
e)          = Expr b -> [Var]
go Expr b
e
    go (Cast Expr b
e CoercionR
_)          = Expr b -> [Var]
go Expr b
e
    go (Case Expr b
e b
_ Type
_ [Alt b]
cs)     = Expr b -> [Var]
go Expr b
e forall a. [a] -> [a] -> [a]
++ forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall a. CBVisitable a => a -> [Var]
readVars [Alt b]
cs
    go Expr b
_                   = []

exprLetVars :: Expr Var -> [Var]
exprLetVars :: Expr Var -> [Var]
exprLetVars = Expr Var -> [Var]
go
  where
    go :: Expr Var -> [Var]
go (Var Var
_)             = []
    go (App Expr Var
e Expr Var
a)           = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Expr Var -> [Var]
go [Expr Var
e, Expr Var
a]
    go (Lam Var
x Expr Var
e)           = Var
x forall a. a -> [a] -> [a]
: Expr Var -> [Var]
go Expr Var
e
    go (Let CoreBind
b Expr Var
e)           = forall a. CBVisitable a => a -> [Var]
letVars CoreBind
b forall a. [a] -> [a] -> [a]
++ Expr Var -> [Var]
go Expr Var
e
    go (Tick CoreTickish
_ Expr Var
e)          = Expr Var -> [Var]
go Expr Var
e
    go (Cast Expr Var
e CoercionR
_)          = Expr Var -> [Var]
go Expr Var
e
    go (Case Expr Var
e Var
x Type
_ [Alt Var]
cs)     = Var
x forall a. a -> [a] -> [a]
: Expr Var -> [Var]
go Expr Var
e forall a. [a] -> [a] -> [a]
++ forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall a. CBVisitable a => a -> [Var]
letVars [Alt Var]
cs
    go Expr Var
_                   = []

exprLiterals :: (CBVisitable (Alt t), CBVisitable (Bind t))
             => Expr t -> [Literal]
exprLiterals :: forall t.
(CBVisitable (Alt t), CBVisitable (Bind t)) =>
Expr t -> [Literal]
exprLiterals = forall {b}.
(CBVisitable (Bind b), CBVisitable (Alt b)) =>
Expr b -> [Literal]
go
  where
    go :: Expr b -> [Literal]
go (Lit Literal
l)             = [Literal
l]
    go (App Expr b
e Expr b
a)           = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Expr b -> [Literal]
go [Expr b
e, Expr b
a]
    go (Let Bind b
b Expr b
e)           = forall a. CBVisitable a => a -> [Literal]
literals Bind b
b forall a. [a] -> [a] -> [a]
++ Expr b -> [Literal]
go Expr b
e
    go (Lam b
_ Expr b
e)           = Expr b -> [Literal]
go Expr b
e
    go (Tick CoreTickish
_ Expr b
e)          = Expr b -> [Literal]
go Expr b
e
    go (Cast Expr b
e CoercionR
_)          = Expr b -> [Literal]
go Expr b
e
    go (Case Expr b
e b
_ Type
_ [Alt b]
cs)     = Expr b -> [Literal]
go Expr b
e forall a. [a] -> [a] -> [a]
++ forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall a. CBVisitable a => a -> [Literal]
literals [Alt b]
cs
    go (Type Type
t)            = Type -> [Literal]
go' Type
t
    go Expr b
_                   = []

    go' :: Type -> [Literal]
go' (LitTy TyLit
tl)         = [TyLit -> Literal
tyLitToLit TyLit
tl]
    go' Type
_                  = []


    tyLitToLit :: TyLit -> Literal
tyLitToLit (CharTyLit Char
c) = Char -> Literal
LitChar Char
c
    tyLitToLit (StrTyLit FastString
fs) = ByteString -> Literal
LitString (FastString -> ByteString
bytesFS FastString
fs)
    tyLitToLit (NumTyLit Integer
i)  = LitNumType -> Integer -> Literal
LitNumber LitNumType
LitNumInt (forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
i)



instance CBVisitable (Alt Var) where
  freeVars :: HashSet Var -> Alt Var -> [Var]
freeVars HashSet Var
env (Alt AltCon
a [Var]
xs Expr Var
e) = forall a. CBVisitable a => HashSet Var -> a -> [Var]
freeVars HashSet Var
env AltCon
a forall a. [a] -> [a] -> [a]
++ forall a. CBVisitable a => HashSet Var -> a -> [Var]
freeVars (forall a. (Eq a, Hashable a) => HashSet a -> [a] -> HashSet a
extendEnv HashSet Var
env [Var]
xs) Expr Var
e
  readVars :: Alt Var -> [Var]
readVars (Alt AltCon
_ [Var]
_ Expr Var
e)       = forall a. CBVisitable a => a -> [Var]
readVars Expr Var
e
  letVars :: Alt Var -> [Var]
letVars  (Alt AltCon
_ [Var]
xs Expr Var
e)       = [Var]
xs forall a. [a] -> [a] -> [a]
++ forall a. CBVisitable a => a -> [Var]
letVars Expr Var
e
  literals :: Alt Var -> [Literal]
literals (Alt AltCon
c [Var]
_ Expr Var
e)       = forall a. CBVisitable a => a -> [Literal]
literals AltCon
c forall a. [a] -> [a] -> [a]
++ forall a. CBVisitable a => a -> [Literal]
literals Expr Var
e

instance CBVisitable AltCon where
  freeVars :: HashSet Var -> AltCon -> [Var]
freeVars HashSet Var
_ (DataAlt DataCon
dc) = [ Var
x | AnId Var
x <- DataCon -> [TyThing]
dataConImplicitTyThings DataCon
dc]
  freeVars HashSet Var
_ AltCon
_            = []
  readVars :: AltCon -> [Var]
readVars AltCon
_              = []
  letVars :: AltCon -> [Var]
letVars  AltCon
_              = []
  literals :: AltCon -> [Literal]
literals (LitAlt Literal
l)     = [Literal
l]
  literals AltCon
_              = []

extendEnv :: (Eq a, Hashable a) => S.HashSet a -> [a] -> S.HashSet a
extendEnv :: forall a. (Eq a, Hashable a) => HashSet a -> [a] -> HashSet a
extendEnv = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a. (Eq a, Hashable a) => a -> HashSet a -> HashSet a
S.insert)

bindings :: Bind t -> [t]
bindings :: forall t. Bind t -> [t]
bindings (NonRec t
x Expr t
_) = [t
x]
bindings (Rec  [(t, Expr t)]
xes  ) = forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> a
fst [(t, Expr t)]
xes

----------------------------------------------------------------------------------------
-- | @BindVisitor@ allows for generic, context sensitive traversals over the @CoreBinds@ 
----------------------------------------------------------------------------------------
data CoreVisitor env acc = CoreVisitor
  { forall env acc. CoreVisitor env acc -> env -> Var -> env
envF  :: env -> Var             -> env
  , forall env acc. CoreVisitor env acc -> env -> acc -> Var -> acc
bindF :: env -> acc -> Var      -> acc
  , forall env acc.
CoreVisitor env acc -> env -> acc -> Expr Var -> acc
exprF :: env -> acc -> CoreExpr -> acc
  }

coreVisitor :: CoreVisitor env acc -> env -> acc -> [CoreBind] -> acc
coreVisitor :: forall env acc.
CoreVisitor env acc -> env -> acc -> [CoreBind] -> acc
coreVisitor CoreVisitor env acc
vis env
cenv acc
cacc [CoreBind]
cbs = forall a b. (a, b) -> b
snd (forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (env, acc) -> CoreBind -> (env, acc)
step (env
cenv, acc
cacc) [CoreBind]
cbs)
  where
    stepXE :: (env, acc) -> (Var, Expr Var) -> (env, acc)
stepXE (env
env, acc
acc) (Var
x,Expr Var
e)   = (env
env', env -> acc -> Expr Var -> acc
stepE env
env' acc
acc'   Expr Var
e)
      where
        env' :: env
env'                  = forall env acc. CoreVisitor env acc -> env -> Var -> env
envF  CoreVisitor env acc
vis env
env     Var
x
        acc' :: acc
acc'                  = forall env acc. CoreVisitor env acc -> env -> acc -> Var -> acc
bindF CoreVisitor env acc
vis env
env acc
acc Var
x

    step :: (env, acc) -> CoreBind -> (env, acc)
step (env, acc)
ea (NonRec Var
x Expr Var
e)      = (env, acc) -> (Var, Expr Var) -> (env, acc)
stepXE (env, acc)
ea (Var
x, Expr Var
e)
    step (env, acc)
ea (Rec    [(Var, Expr Var)]
xes)      = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (env, acc) -> (Var, Expr Var) -> (env, acc)
stepXE (env, acc)
ea [(Var, Expr Var)]
xes

    -- step (env, acc) (NonRec x e) = stepXE env acc x e 
    -- step (env, acc) (Rec    xes) = (env', foldl' (stepE env') acc' es) 
      -- where 
        -- acc'                     = foldl' (bindF vis env') acc xs
        -- env'                     = foldl' (envF  vis)      env xs 
        -- xs                       = fst <$> xes 
        -- es                       = snd <$> xes
        -- foldl' (\(env, acc) (x, e) ->  )

    stepE :: env -> acc -> Expr Var -> acc
stepE env
env acc
acc Expr Var
e              = env -> acc -> Expr Var -> acc
goE env
env (forall env acc.
CoreVisitor env acc -> env -> acc -> Expr Var -> acc
exprF CoreVisitor env acc
vis env
env acc
acc Expr Var
e) Expr Var
e

    goE :: env -> acc -> Expr Var -> acc
goE env
_   acc
acc (Var Var
_)          = acc
acc
    goE env
env acc
acc (App Expr Var
e1 Expr Var
e2)      = env -> acc -> Expr Var -> acc
stepE  env
env (env -> acc -> Expr Var -> acc
stepE env
env acc
acc Expr Var
e1) Expr Var
e2
    goE env
env acc
acc (Tick CoreTickish
_ Expr Var
e)       = env -> acc -> Expr Var -> acc
stepE  env
env acc
acc Expr Var
e
    goE env
env acc
acc (Cast Expr Var
e CoercionR
_)       = env -> acc -> Expr Var -> acc
stepE  env
env acc
acc Expr Var
e
    goE env
env acc
acc (Lam Var
x Expr Var
e)        = forall a b. (a, b) -> b
snd ((env, acc) -> (Var, Expr Var) -> (env, acc)
stepXE (env
env, acc
acc) (Var
x, Expr Var
e))
    goE env
env acc
acc (Let CoreBind
b Expr Var
e)        = env -> acc -> Expr Var -> acc
stepE env
env' acc
acc' Expr Var
e where (env
env', acc
acc') = (env, acc) -> CoreBind -> (env, acc)
step (env
env, acc
acc) CoreBind
b
    goE env
env acc
acc (Case Expr Var
e Var
_ Type
_ [Alt Var]
cs)  = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (env -> acc -> Alt Var -> acc
goC env
env) (env -> acc -> Expr Var -> acc
stepE env
env acc
acc Expr Var
e) [Alt Var]
cs
    goE env
_   acc
acc Expr Var
_                = acc
acc

    goC :: env -> acc -> Alt Var -> acc
goC env
env acc
acc (Alt AltCon
_ [Var]
xs Expr Var
e)     = env -> acc -> Expr Var -> acc
stepE  env
env' acc
acc' Expr Var
e
      where
        env' :: env
env'                     = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (forall env acc. CoreVisitor env acc -> env -> Var -> env
envF  CoreVisitor env acc
vis)     env
env [Var]
xs
        acc' :: acc
acc'                     = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (forall env acc. CoreVisitor env acc -> env -> acc -> Var -> acc
bindF CoreVisitor env acc
vis env
env) acc
acc [Var]
xs