{-# LANGUAGE CPP #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE PatternGuards #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE BangPatterns  #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}

module Language.Fixpoint.Types.Visitor (
  -- * Visitor
     Visitor (..)
  ,  Visitable (..)

  -- * Extracting Symbolic Constants (String Literals)
  ,  SymConsts (..)

  -- * Default Visitor
  , defaultVisitor

  -- * Transformers
  , trans

  -- * Accumulators
  , fold

  -- * Clients
  , stripCasts
  , kvars, eapps
  , size, lamSize
  , envKVars
  , envKVarsN
  , rhsKVars
  , mapKVars, mapKVars', mapGVars', mapKVarSubsts
  , mapExpr, mapMExpr

  -- * Coercion Substitutions
  , CoSub
  , applyCoSub

  -- * Predicates on Constraints
  , isConcC , isKvarC

  -- * Sorts
  , foldSort
  , mapSort
  , foldDataDecl

  , (<$$>)


  ) where

-- import           Control.Monad.Trans.State.Strict (State, modify, runState)
-- import           Control.DeepSeq
#if !MIN_VERSION_base(4,14,0)
import           Data.Semigroup      (Semigroup (..))
#endif

import           Control.Monad.State.Strict
import qualified Data.HashSet        as S
import qualified Data.HashMap.Strict as M
import qualified Data.List           as L
import           Language.Fixpoint.Types hiding (mapSort)
import qualified Language.Fixpoint.Misc as Misc




data Visitor acc ctx = Visitor {
 -- | Context @ctx@ is built in a "top-down" fashion; not "across" siblings
    Visitor acc ctx -> ctx -> Expr -> ctx
ctxExpr :: ctx -> Expr -> ctx

  -- | Transforms can access current @ctx@
  , Visitor acc ctx -> ctx -> Expr -> Expr
txExpr  :: ctx -> Expr -> Expr

  -- | Accumulations can access current @ctx@; @acc@ value is monoidal
  , Visitor acc ctx -> ctx -> Expr -> acc
accExpr :: ctx -> Expr -> acc
  }

---------------------------------------------------------------------------------
defaultVisitor :: (Monoid acc) => Visitor acc ctx
---------------------------------------------------------------------------------
defaultVisitor :: Visitor acc ctx
defaultVisitor = Visitor :: forall acc ctx.
(ctx -> Expr -> ctx)
-> (ctx -> Expr -> Expr) -> (ctx -> Expr -> acc) -> Visitor acc ctx
Visitor
  { ctxExpr :: ctx -> Expr -> ctx
ctxExpr    = ctx -> Expr -> ctx
forall a b. a -> b -> a
const
  , txExpr :: ctx -> Expr -> Expr
txExpr     = \ctx
_ Expr
x -> Expr
x
  , accExpr :: ctx -> Expr -> acc
accExpr    = \ctx
_ Expr
_ -> acc
forall a. Monoid a => a
mempty
  }

------------------------------------------------------------------------

fold         :: (Visitable t, Monoid a) => Visitor a ctx -> ctx -> a -> t -> a
fold :: Visitor a ctx -> ctx -> a -> t -> a
fold Visitor a ctx
v ctx
c a
a t
t = (t, a) -> a
forall a b. (a, b) -> b
snd ((t, a) -> a) -> (t, a) -> a
forall a b. (a -> b) -> a -> b
$ Visitor a ctx
-> ctx
-> a
-> (Visitor a ctx -> ctx -> t -> State a t)
-> t
-> (t, a)
forall a ctx t.
Visitor a ctx
-> ctx
-> a
-> (Visitor a ctx -> ctx -> t -> State a t)
-> t
-> (t, a)
execVisitM Visitor a ctx
v ctx
c a
a Visitor a ctx -> ctx -> t -> State a t
forall t a c.
(Visitable t, Monoid a) =>
Visitor a c -> c -> t -> VisitM a t
visit t
t

trans        :: (Visitable t, Monoid a) => Visitor a ctx -> ctx -> a -> t -> t
trans :: Visitor a ctx -> ctx -> a -> t -> t
trans Visitor a ctx
v ctx
c a
_ t
z = (t, a) -> t
forall a b. (a, b) -> a
fst ((t, a) -> t) -> (t, a) -> t
forall a b. (a -> b) -> a -> b
$ Visitor a ctx
-> ctx
-> a
-> (Visitor a ctx -> ctx -> t -> State a t)
-> t
-> (t, a)
forall a ctx t.
Visitor a ctx
-> ctx
-> a
-> (Visitor a ctx -> ctx -> t -> State a t)
-> t
-> (t, a)
execVisitM Visitor a ctx
v ctx
c a
forall a. Monoid a => a
mempty Visitor a ctx -> ctx -> t -> State a t
forall t a c.
(Visitable t, Monoid a) =>
Visitor a c -> c -> t -> VisitM a t
visit t
z

execVisitM :: Visitor a ctx -> ctx -> a -> (Visitor a ctx -> ctx -> t -> State a t) -> t -> (t, a)
execVisitM :: Visitor a ctx
-> ctx
-> a
-> (Visitor a ctx -> ctx -> t -> State a t)
-> t
-> (t, a)
execVisitM Visitor a ctx
v ctx
c a
a Visitor a ctx -> ctx -> t -> State a t
f t
x = State a t -> a -> (t, a)
forall s a. State s a -> s -> (a, s)
runState (Visitor a ctx -> ctx -> t -> State a t
f Visitor a ctx
v ctx
c t
x) a
a

type VisitM acc = State acc

accum :: (Monoid a) => a -> VisitM a ()
accum :: a -> VisitM a ()
accum !a
z = (a -> a) -> VisitM a ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (a -> a -> a
forall a. Monoid a => a -> a -> a
mappend a
z)
  -- do
  -- !cur <- get
  -- put ((mappend $!! z) $!! cur)

(<$$>) :: (Monad m) => (a -> m b) -> [a] -> m [b]
a -> m b
f <$$> :: (a -> m b) -> [a] -> m [b]
<$$> [a]
xs = a -> m b
f (a -> m b) -> [a] -> m [b]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
Misc.<$$> [a]
xs

-- (<$$>) ::  (Applicative f) => (a -> f b) -> [a] -> f [b]
-- f <$$> x = traverse f x
-- _ <$$> []     = return []
-- f <$$> (x:xs) = do
  -- !y  <- f x
  -- !ys <- f <$$> xs
  -- return (y:ys)
------------------------------------------------------------------------------
class Visitable t where
  visit :: (Monoid a) => Visitor a c -> c -> t -> VisitM a t

instance Visitable Expr where
  visit :: Visitor a c -> c -> Expr -> VisitM a Expr
visit = Visitor a c -> c -> Expr -> VisitM a Expr
forall a c. Monoid a => Visitor a c -> c -> Expr -> VisitM a Expr
visitExpr

instance Visitable Reft where
  visit :: Visitor a c -> c -> Reft -> VisitM a Reft
visit Visitor a c
v c
c (Reft (Symbol
x, Expr
ra)) = ((Symbol, Expr) -> Reft
Reft ((Symbol, Expr) -> Reft)
-> (Expr -> (Symbol, Expr)) -> Expr -> Reft
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Symbol
x, )) (Expr -> Reft) -> StateT a Identity Expr -> VisitM a Reft
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Visitor a c -> c -> Expr -> StateT a Identity Expr
forall t a c.
(Visitable t, Monoid a) =>
Visitor a c -> c -> t -> VisitM a t
visit Visitor a c
v c
c Expr
ra

instance Visitable SortedReft where
  visit :: Visitor a c -> c -> SortedReft -> VisitM a SortedReft
visit Visitor a c
v c
c (RR Sort
t Reft
r) = Sort -> Reft -> SortedReft
RR Sort
t (Reft -> SortedReft)
-> StateT a Identity Reft -> VisitM a SortedReft
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Visitor a c -> c -> Reft -> StateT a Identity Reft
forall t a c.
(Visitable t, Monoid a) =>
Visitor a c -> c -> t -> VisitM a t
visit Visitor a c
v c
c Reft
r

instance Visitable (Symbol, SortedReft) where
  visit :: Visitor a c
-> c -> (Symbol, SortedReft) -> VisitM a (Symbol, SortedReft)
visit Visitor a c
v c
c (Symbol
sym, SortedReft
sr) = (Symbol
sym, ) (SortedReft -> (Symbol, SortedReft))
-> StateT a Identity SortedReft -> VisitM a (Symbol, SortedReft)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Visitor a c -> c -> SortedReft -> StateT a Identity SortedReft
forall t a c.
(Visitable t, Monoid a) =>
Visitor a c -> c -> t -> VisitM a t
visit Visitor a c
v c
c SortedReft
sr

instance Visitable BindEnv where
  visit :: Visitor a c -> c -> BindEnv -> VisitM a BindEnv
visit Visitor a c
v c
c = ((Symbol, SortedReft) -> StateT a Identity (Symbol, SortedReft))
-> BindEnv -> VisitM a BindEnv
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Visitor a c
-> c
-> (Symbol, SortedReft)
-> StateT a Identity (Symbol, SortedReft)
forall t a c.
(Visitable t, Monoid a) =>
Visitor a c -> c -> t -> VisitM a t
visit Visitor a c
v c
c)

---------------------------------------------------------------------------------
-- WARNING: these instances were written for mapKVars over GInfos only;
-- check that they behave as expected before using with other clients.
instance Visitable (SimpC a) where
  visit :: Visitor a c -> c -> SimpC a -> VisitM a (SimpC a)
visit Visitor a c
v c
c SimpC a
x = do
    Expr
rhs' <- Visitor a c -> c -> Expr -> VisitM a Expr
forall t a c.
(Visitable t, Monoid a) =>
Visitor a c -> c -> t -> VisitM a t
visit Visitor a c
v c
c (SimpC a -> Expr
forall a. SimpC a -> Expr
_crhs SimpC a
x)
    SimpC a -> VisitM a (SimpC a)
forall (m :: * -> *) a. Monad m => a -> m a
return SimpC a
x { _crhs :: Expr
_crhs = Expr
rhs' }

instance Visitable (SubC a) where
  visit :: Visitor a c -> c -> SubC a -> VisitM a (SubC a)
visit Visitor a c
v c
c SubC a
x = do
    SortedReft
lhs' <- Visitor a c -> c -> SortedReft -> VisitM a SortedReft
forall t a c.
(Visitable t, Monoid a) =>
Visitor a c -> c -> t -> VisitM a t
visit Visitor a c
v c
c (SubC a -> SortedReft
forall a. SubC a -> SortedReft
slhs SubC a
x)
    SortedReft
rhs' <- Visitor a c -> c -> SortedReft -> VisitM a SortedReft
forall t a c.
(Visitable t, Monoid a) =>
Visitor a c -> c -> t -> VisitM a t
visit Visitor a c
v c
c (SubC a -> SortedReft
forall a. SubC a -> SortedReft
srhs SubC a
x)
    SubC a -> VisitM a (SubC a)
forall (m :: * -> *) a. Monad m => a -> m a
return SubC a
x { slhs :: SortedReft
slhs = SortedReft
lhs', srhs :: SortedReft
srhs = SortedReft
rhs' }

instance (Visitable (c a)) => Visitable (GInfo c a) where
  visit :: Visitor a c -> c -> GInfo c a -> VisitM a (GInfo c a)
visit Visitor a c
v c
c GInfo c a
x = do
    HashMap SubcId (c a)
cm' <- (c a -> StateT a Identity (c a))
-> HashMap SubcId (c a) -> StateT a Identity (HashMap SubcId (c a))
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Visitor a c -> c -> c a -> StateT a Identity (c a)
forall t a c.
(Visitable t, Monoid a) =>
Visitor a c -> c -> t -> VisitM a t
visit Visitor a c
v c
c) (GInfo c a -> HashMap SubcId (c a)
forall (c :: * -> *) a. GInfo c a -> HashMap SubcId (c a)
cm GInfo c a
x)
    BindEnv
bs' <- Visitor a c -> c -> BindEnv -> VisitM a BindEnv
forall t a c.
(Visitable t, Monoid a) =>
Visitor a c -> c -> t -> VisitM a t
visit Visitor a c
v c
c (GInfo c a -> BindEnv
forall (c :: * -> *) a. GInfo c a -> BindEnv
bs GInfo c a
x)
    AxiomEnv
ae' <- Visitor a c -> c -> AxiomEnv -> VisitM a AxiomEnv
forall t a c.
(Visitable t, Monoid a) =>
Visitor a c -> c -> t -> VisitM a t
visit Visitor a c
v c
c (GInfo c a -> AxiomEnv
forall (c :: * -> *) a. GInfo c a -> AxiomEnv
ae GInfo c a
x) 
    GInfo c a -> VisitM a (GInfo c a)
forall (m :: * -> *) a. Monad m => a -> m a
return GInfo c a
x { cm :: HashMap SubcId (c a)
cm = HashMap SubcId (c a)
cm', bs :: BindEnv
bs = BindEnv
bs', ae :: AxiomEnv
ae = AxiomEnv
ae' }

instance Visitable AxiomEnv where 
  visit :: Visitor a c -> c -> AxiomEnv -> VisitM a AxiomEnv
visit Visitor a c
v c
c AxiomEnv
x = do 
    [Equation]
eqs'    <- (Equation -> StateT a Identity Equation)
-> [Equation] -> StateT a Identity [Equation]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Visitor a c -> c -> Equation -> StateT a Identity Equation
forall t a c.
(Visitable t, Monoid a) =>
Visitor a c -> c -> t -> VisitM a t
visit Visitor a c
v c
c) (AxiomEnv -> [Equation]
aenvEqs   AxiomEnv
x) 
    [Rewrite]
simpls' <- (Rewrite -> StateT a Identity Rewrite)
-> [Rewrite] -> StateT a Identity [Rewrite]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Visitor a c -> c -> Rewrite -> StateT a Identity Rewrite
forall t a c.
(Visitable t, Monoid a) =>
Visitor a c -> c -> t -> VisitM a t
visit Visitor a c
v c
c) (AxiomEnv -> [Rewrite]
aenvSimpl AxiomEnv
x) 
    AxiomEnv -> VisitM a AxiomEnv
forall (m :: * -> *) a. Monad m => a -> m a
return AxiomEnv
x { aenvEqs :: [Equation]
aenvEqs = [Equation]
eqs' , aenvSimpl :: [Rewrite]
aenvSimpl = [Rewrite]
simpls'} 

instance Visitable Equation where 
  visit :: Visitor a c -> c -> Equation -> VisitM a Equation
visit Visitor a c
v c
c Equation
eq = do 
    Expr
body' <- Visitor a c -> c -> Expr -> VisitM a Expr
forall t a c.
(Visitable t, Monoid a) =>
Visitor a c -> c -> t -> VisitM a t
visit Visitor a c
v c
c (Equation -> Expr
eqBody Equation
eq) 
    Equation -> VisitM a Equation
forall (m :: * -> *) a. Monad m => a -> m a
return Equation
eq { eqBody :: Expr
eqBody = Expr
body' } 

instance Visitable Rewrite where 
  visit :: Visitor a c -> c -> Rewrite -> VisitM a Rewrite
visit Visitor a c
v c
c Rewrite
rw = do 
    Expr
body' <- Visitor a c -> c -> Expr -> VisitM a Expr
forall t a c.
(Visitable t, Monoid a) =>
Visitor a c -> c -> t -> VisitM a t
visit Visitor a c
v c
c (Rewrite -> Expr
smBody Rewrite
rw) 
    Rewrite -> VisitM a Rewrite
forall (m :: * -> *) a. Monad m => a -> m a
return Rewrite
rw { smBody :: Expr
smBody = Expr
body' } 

---------------------------------------------------------------------------------
visitExpr :: (Monoid a) => Visitor a ctx -> ctx -> Expr -> VisitM a Expr
visitExpr :: Visitor a ctx -> ctx -> Expr -> VisitM a Expr
visitExpr !Visitor a ctx
v    = ctx -> Expr -> VisitM a Expr
vE
  where
    vE :: ctx -> Expr -> VisitM a Expr
vE !ctx
c !Expr
e    = do {-# SCC "visitExpr.vE.accum" #-} a -> VisitM a ()
forall a. Monoid a => a -> VisitM a ()
accum a
acc
                     {-# SCC "visitExpr.vE.step" #-}  ctx -> Expr -> VisitM a Expr
step ctx
c' Expr
e'
      where !c' :: ctx
c'  = Visitor a ctx -> ctx -> Expr -> ctx
forall acc ctx. Visitor acc ctx -> ctx -> Expr -> ctx
ctxExpr Visitor a ctx
v ctx
c  Expr
e
            !e' :: Expr
e'  = Visitor a ctx -> ctx -> Expr -> Expr
forall acc ctx. Visitor acc ctx -> ctx -> Expr -> Expr
txExpr  Visitor a ctx
v ctx
c' Expr
e
            !acc :: a
acc = Visitor a ctx -> ctx -> Expr -> a
forall acc ctx. Visitor acc ctx -> ctx -> Expr -> acc
accExpr Visitor a ctx
v ctx
c' Expr
e
    step :: ctx -> Expr -> VisitM a Expr
step ctx
_ !e :: Expr
e@(ESym SymConst
_)       = Expr -> VisitM a Expr
forall (m :: * -> *) a. Monad m => a -> m a
return Expr
e
    step ctx
_ !e :: Expr
e@(ECon Constant
_)       = Expr -> VisitM a Expr
forall (m :: * -> *) a. Monad m => a -> m a
return Expr
e
    step ctx
_ !e :: Expr
e@(EVar Symbol
_)       = Expr -> VisitM a Expr
forall (m :: * -> *) a. Monad m => a -> m a
return Expr
e
    step !ctx
c !(EApp Expr
f Expr
e)      = Expr -> Expr -> Expr
EApp        (Expr -> Expr -> Expr)
-> VisitM a Expr -> StateT a Identity (Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ctx -> Expr -> VisitM a Expr
vE ctx
c Expr
f  StateT a Identity (Expr -> Expr) -> VisitM a Expr -> VisitM a Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ctx -> Expr -> VisitM a Expr
vE ctx
c Expr
e
    step !ctx
c !(ENeg Expr
e)        = Expr -> Expr
ENeg        (Expr -> Expr) -> VisitM a Expr -> VisitM a Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ctx -> Expr -> VisitM a Expr
vE ctx
c Expr
e
    step !ctx
c !(EBin Bop
o Expr
e1 Expr
e2)  = Bop -> Expr -> Expr -> Expr
EBin Bop
o      (Expr -> Expr -> Expr)
-> VisitM a Expr -> StateT a Identity (Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ctx -> Expr -> VisitM a Expr
vE ctx
c Expr
e1 StateT a Identity (Expr -> Expr) -> VisitM a Expr -> VisitM a Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ctx -> Expr -> VisitM a Expr
vE ctx
c Expr
e2
    step !ctx
c !(EIte Expr
p Expr
e1 Expr
e2)  = Expr -> Expr -> Expr -> Expr
EIte        (Expr -> Expr -> Expr -> Expr)
-> VisitM a Expr -> StateT a Identity (Expr -> Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ctx -> Expr -> VisitM a Expr
vE ctx
c Expr
p  StateT a Identity (Expr -> Expr -> Expr)
-> VisitM a Expr -> StateT a Identity (Expr -> Expr)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ctx -> Expr -> VisitM a Expr
vE ctx
c Expr
e1 StateT a Identity (Expr -> Expr) -> VisitM a Expr -> VisitM a Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ctx -> Expr -> VisitM a Expr
vE ctx
c Expr
e2
    step !ctx
c !(ECst Expr
e Sort
t)      = (Expr -> Sort -> Expr
`ECst` Sort
t)  (Expr -> Expr) -> VisitM a Expr -> VisitM a Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ctx -> Expr -> VisitM a Expr
vE ctx
c Expr
e
    step !ctx
c !(PAnd  [Expr]
ps)      = [Expr] -> Expr
PAnd        ([Expr] -> Expr) -> StateT a Identity [Expr] -> VisitM a Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ctx -> Expr -> VisitM a Expr
vE ctx
c (Expr -> VisitM a Expr) -> [Expr] -> StateT a Identity [Expr]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
<$$> [Expr]
ps)
    step !ctx
c !(POr  [Expr]
ps)       = [Expr] -> Expr
POr         ([Expr] -> Expr) -> StateT a Identity [Expr] -> VisitM a Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ctx -> Expr -> VisitM a Expr
vE ctx
c (Expr -> VisitM a Expr) -> [Expr] -> StateT a Identity [Expr]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
<$$> [Expr]
ps)
    step !ctx
c !(PNot Expr
p)        = Expr -> Expr
PNot        (Expr -> Expr) -> VisitM a Expr -> VisitM a Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ctx -> Expr -> VisitM a Expr
vE ctx
c Expr
p
    step !ctx
c !(PImp Expr
p1 Expr
p2)    = Expr -> Expr -> Expr
PImp        (Expr -> Expr -> Expr)
-> VisitM a Expr -> StateT a Identity (Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ctx -> Expr -> VisitM a Expr
vE ctx
c Expr
p1 StateT a Identity (Expr -> Expr) -> VisitM a Expr -> VisitM a Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ctx -> Expr -> VisitM a Expr
vE ctx
c Expr
p2
    step !ctx
c !(PIff Expr
p1 Expr
p2)    = Expr -> Expr -> Expr
PIff        (Expr -> Expr -> Expr)
-> VisitM a Expr -> StateT a Identity (Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ctx -> Expr -> VisitM a Expr
vE ctx
c Expr
p1 StateT a Identity (Expr -> Expr) -> VisitM a Expr -> VisitM a Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ctx -> Expr -> VisitM a Expr
vE ctx
c Expr
p2
    step !ctx
c !(PAtom Brel
r Expr
e1 Expr
e2) = Brel -> Expr -> Expr -> Expr
PAtom Brel
r     (Expr -> Expr -> Expr)
-> VisitM a Expr -> StateT a Identity (Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ctx -> Expr -> VisitM a Expr
vE ctx
c Expr
e1 StateT a Identity (Expr -> Expr) -> VisitM a Expr -> VisitM a Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> ctx -> Expr -> VisitM a Expr
vE ctx
c Expr
e2
    step !ctx
c !(PAll [(Symbol, Sort)]
xts Expr
p)    = [(Symbol, Sort)] -> Expr -> Expr
PAll   [(Symbol, Sort)]
xts  (Expr -> Expr) -> VisitM a Expr -> VisitM a Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ctx -> Expr -> VisitM a Expr
vE ctx
c Expr
p
    step !ctx
c !(ELam (Symbol
x,Sort
t) Expr
e)  = (Symbol, Sort) -> Expr -> Expr
ELam (Symbol
x,Sort
t)  (Expr -> Expr) -> VisitM a Expr -> VisitM a Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ctx -> Expr -> VisitM a Expr
vE ctx
c Expr
e
    step !ctx
c !(ECoerc Sort
a Sort
t Expr
e)  = Sort -> Sort -> Expr -> Expr
ECoerc Sort
a Sort
t  (Expr -> Expr) -> VisitM a Expr -> VisitM a Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ctx -> Expr -> VisitM a Expr
vE ctx
c Expr
e
    step !ctx
c !(PExist [(Symbol, Sort)]
xts Expr
p)  = [(Symbol, Sort)] -> Expr -> Expr
PExist [(Symbol, Sort)]
xts  (Expr -> Expr) -> VisitM a Expr -> VisitM a Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ctx -> Expr -> VisitM a Expr
vE ctx
c Expr
p
    step !ctx
c !(ETApp Expr
e Sort
s)     = (Expr -> Sort -> Expr
`ETApp` Sort
s) (Expr -> Expr) -> VisitM a Expr -> VisitM a Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ctx -> Expr -> VisitM a Expr
vE ctx
c Expr
e
    step !ctx
c !(ETAbs Expr
e Symbol
s)     = (Expr -> Symbol -> Expr
`ETAbs` Symbol
s) (Expr -> Expr) -> VisitM a Expr -> VisitM a Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ctx -> Expr -> VisitM a Expr
vE ctx
c Expr
e
    step ctx
_  !p :: Expr
p@(PKVar KVar
_ Subst
_)   = Expr -> VisitM a Expr
forall (m :: * -> *) a. Monad m => a -> m a
return Expr
p
    step !ctx
c !(PGrad KVar
k Subst
su GradInfo
i Expr
e) = KVar -> Subst -> GradInfo -> Expr -> Expr
PGrad KVar
k Subst
su GradInfo
i (Expr -> Expr) -> VisitM a Expr -> VisitM a Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ctx -> Expr -> VisitM a Expr
vE ctx
c Expr
e

mapKVars :: Visitable t => (KVar -> Maybe Expr) -> t -> t
mapKVars :: (KVar -> Maybe Expr) -> t -> t
mapKVars KVar -> Maybe Expr
f = ((KVar, Subst) -> Maybe Expr) -> t -> t
forall t. Visitable t => ((KVar, Subst) -> Maybe Expr) -> t -> t
mapKVars' (KVar, Subst) -> Maybe Expr
forall b. (KVar, b) -> Maybe Expr
f'
  where
    f' :: (KVar, b) -> Maybe Expr
f' (KVar
kv', b
_) = KVar -> Maybe Expr
f KVar
kv'

mapKVars' :: Visitable t => ((KVar, Subst) -> Maybe Expr) -> t -> t
mapKVars' :: ((KVar, Subst) -> Maybe Expr) -> t -> t
mapKVars' (KVar, Subst) -> Maybe Expr
f            = Visitor () () -> () -> () -> t -> t
forall t a ctx.
(Visitable t, Monoid a) =>
Visitor a ctx -> ctx -> a -> t -> t
trans Visitor () ()
forall p. Visitor () p
kvVis () ()
  where
    kvVis :: Visitor () p
kvVis              = Visitor () p
forall acc ctx. Monoid acc => Visitor acc ctx
defaultVisitor { txExpr :: p -> Expr -> Expr
txExpr = p -> Expr -> Expr
forall p. p -> Expr -> Expr
txK }
    txK :: p -> Expr -> Expr
txK p
_ (PKVar KVar
k Subst
su)
      | Just Expr
p' <- (KVar, Subst) -> Maybe Expr
f (KVar
k, Subst
su) = Subst -> Expr -> Expr
forall a. Subable a => Subst -> a -> a
subst Subst
su Expr
p'
    txK p
_ (PGrad KVar
k Subst
su GradInfo
_ Expr
_)
      | Just Expr
p' <- (KVar, Subst) -> Maybe Expr
f (KVar
k, Subst
su) = Subst -> Expr -> Expr
forall a. Subable a => Subst -> a -> a
subst Subst
su Expr
p'
    txK p
_ Expr
p            = Expr
p



mapGVars' :: Visitable t => ((KVar, Subst) -> Maybe Expr) -> t -> t
mapGVars' :: ((KVar, Subst) -> Maybe Expr) -> t -> t
mapGVars' (KVar, Subst) -> Maybe Expr
f            = Visitor () () -> () -> () -> t -> t
forall t a ctx.
(Visitable t, Monoid a) =>
Visitor a ctx -> ctx -> a -> t -> t
trans Visitor () ()
forall p. Visitor () p
kvVis () ()
  where
    kvVis :: Visitor () p
kvVis              = Visitor () p
forall acc ctx. Monoid acc => Visitor acc ctx
defaultVisitor { txExpr :: p -> Expr -> Expr
txExpr = p -> Expr -> Expr
forall p. p -> Expr -> Expr
txK }
    txK :: p -> Expr -> Expr
txK p
_ (PGrad KVar
k Subst
su GradInfo
_ Expr
_)
      | Just Expr
p' <- (KVar, Subst) -> Maybe Expr
f (KVar
k, Subst
su) = Subst -> Expr -> Expr
forall a. Subable a => Subst -> a -> a
subst Subst
su Expr
p'
    txK p
_ Expr
p            = Expr
p

mapExpr :: Visitable t => (Expr -> Expr) -> t -> t
mapExpr :: (Expr -> Expr) -> t -> t
mapExpr Expr -> Expr
f = Visitor () () -> () -> () -> t -> t
forall t a ctx.
(Visitable t, Monoid a) =>
Visitor a ctx -> ctx -> a -> t -> t
trans (Visitor () ()
forall acc ctx. Monoid acc => Visitor acc ctx
defaultVisitor {txExpr :: () -> Expr -> Expr
txExpr = (Expr -> Expr) -> () -> Expr -> Expr
forall a b. a -> b -> a
const Expr -> Expr
f}) () ()


mapMExpr :: (Monad m) => (Expr -> m Expr) -> Expr -> m Expr
mapMExpr :: (Expr -> m Expr) -> Expr -> m Expr
mapMExpr Expr -> m Expr
f = Expr -> m Expr
go
  where
    go :: Expr -> m Expr
go e :: Expr
e@(ESym SymConst
_)      = Expr -> m Expr
f Expr
e
    go e :: Expr
e@(ECon Constant
_)      = Expr -> m Expr
f Expr
e
    go e :: Expr
e@(EVar Symbol
_)      = Expr -> m Expr
f Expr
e
    go e :: Expr
e@(PKVar KVar
_ Subst
_)   = Expr -> m Expr
f Expr
e
    go (PGrad KVar
k Subst
s GradInfo
i Expr
e) = Expr -> m Expr
f (Expr -> m Expr) -> m Expr -> m Expr
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (KVar -> Subst -> GradInfo -> Expr -> Expr
PGrad KVar
k Subst
s GradInfo
i (Expr -> Expr) -> m Expr -> m Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>  Expr -> m Expr
go Expr
e                     )
    go (ENeg Expr
e)        = Expr -> m Expr
f (Expr -> m Expr) -> m Expr -> m Expr
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Expr -> Expr
ENeg        (Expr -> Expr) -> m Expr -> m Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>  Expr -> m Expr
go Expr
e                     )
    go (PNot Expr
p)        = Expr -> m Expr
f (Expr -> m Expr) -> m Expr -> m Expr
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Expr -> Expr
PNot        (Expr -> Expr) -> m Expr -> m Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>  Expr -> m Expr
go Expr
p                     )
    go (ECst Expr
e Sort
t)      = Expr -> m Expr
f (Expr -> m Expr) -> m Expr -> m Expr
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ((Expr -> Sort -> Expr
`ECst` Sort
t)  (Expr -> Expr) -> m Expr -> m Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>  Expr -> m Expr
go Expr
e                     )
    go (PAll [(Symbol, Sort)]
xts Expr
p)    = Expr -> m Expr
f (Expr -> m Expr) -> m Expr -> m Expr
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ([(Symbol, Sort)] -> Expr -> Expr
PAll   [(Symbol, Sort)]
xts  (Expr -> Expr) -> m Expr -> m Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>  Expr -> m Expr
go Expr
p                     )
    go (ELam (Symbol
x,Sort
t) Expr
e)  = Expr -> m Expr
f (Expr -> m Expr) -> m Expr -> m Expr
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ((Symbol, Sort) -> Expr -> Expr
ELam (Symbol
x,Sort
t)  (Expr -> Expr) -> m Expr -> m Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>  Expr -> m Expr
go Expr
e                     )
    go (ECoerc Sort
a Sort
t Expr
e)  = Expr -> m Expr
f (Expr -> m Expr) -> m Expr -> m Expr
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Sort -> Sort -> Expr -> Expr
ECoerc Sort
a Sort
t  (Expr -> Expr) -> m Expr -> m Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>  Expr -> m Expr
go Expr
e                     )
    go (PExist [(Symbol, Sort)]
xts Expr
p)  = Expr -> m Expr
f (Expr -> m Expr) -> m Expr -> m Expr
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ([(Symbol, Sort)] -> Expr -> Expr
PExist [(Symbol, Sort)]
xts  (Expr -> Expr) -> m Expr -> m Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>  Expr -> m Expr
go Expr
p                     )
    go (ETApp Expr
e Sort
s)     = Expr -> m Expr
f (Expr -> m Expr) -> m Expr -> m Expr
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ((Expr -> Sort -> Expr
`ETApp` Sort
s) (Expr -> Expr) -> m Expr -> m Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>  Expr -> m Expr
go Expr
e                     )
    go (ETAbs Expr
e Symbol
s)     = Expr -> m Expr
f (Expr -> m Expr) -> m Expr -> m Expr
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ((Expr -> Symbol -> Expr
`ETAbs` Symbol
s) (Expr -> Expr) -> m Expr -> m Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>  Expr -> m Expr
go Expr
e                     )
    go (EApp Expr
g Expr
e)      = Expr -> m Expr
f (Expr -> m Expr) -> m Expr -> m Expr
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Expr -> Expr -> Expr
EApp        (Expr -> Expr -> Expr) -> m Expr -> m (Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>  Expr -> m Expr
go Expr
g  m (Expr -> Expr) -> m Expr -> m Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr -> m Expr
go Expr
e           )
    go (EBin Bop
o Expr
e1 Expr
e2)  = Expr -> m Expr
f (Expr -> m Expr) -> m Expr -> m Expr
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Bop -> Expr -> Expr -> Expr
EBin Bop
o      (Expr -> Expr -> Expr) -> m Expr -> m (Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>  Expr -> m Expr
go Expr
e1 m (Expr -> Expr) -> m Expr -> m Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr -> m Expr
go Expr
e2          )
    go (PImp Expr
p1 Expr
p2)    = Expr -> m Expr
f (Expr -> m Expr) -> m Expr -> m Expr
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Expr -> Expr -> Expr
PImp        (Expr -> Expr -> Expr) -> m Expr -> m (Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>  Expr -> m Expr
go Expr
p1 m (Expr -> Expr) -> m Expr -> m Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr -> m Expr
go Expr
p2          )
    go (PIff Expr
p1 Expr
p2)    = Expr -> m Expr
f (Expr -> m Expr) -> m Expr -> m Expr
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Expr -> Expr -> Expr
PIff        (Expr -> Expr -> Expr) -> m Expr -> m (Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>  Expr -> m Expr
go Expr
p1 m (Expr -> Expr) -> m Expr -> m Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr -> m Expr
go Expr
p2          )
    go (PAtom Brel
r Expr
e1 Expr
e2) = Expr -> m Expr
f (Expr -> m Expr) -> m Expr -> m Expr
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Brel -> Expr -> Expr -> Expr
PAtom Brel
r     (Expr -> Expr -> Expr) -> m Expr -> m (Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>  Expr -> m Expr
go Expr
e1 m (Expr -> Expr) -> m Expr -> m Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr -> m Expr
go Expr
e2          )
    go (EIte Expr
p Expr
e1 Expr
e2)  = Expr -> m Expr
f (Expr -> m Expr) -> m Expr -> m Expr
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Expr -> Expr -> Expr -> Expr
EIte        (Expr -> Expr -> Expr -> Expr)
-> m Expr -> m (Expr -> Expr -> Expr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>  Expr -> m Expr
go Expr
p  m (Expr -> Expr -> Expr) -> m Expr -> m (Expr -> Expr)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr -> m Expr
go Expr
e1 m (Expr -> Expr) -> m Expr -> m Expr
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Expr -> m Expr
go Expr
e2)
    go (PAnd  [Expr]
ps)      = Expr -> m Expr
f (Expr -> m Expr) -> m Expr -> m Expr
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ([Expr] -> Expr
PAnd        ([Expr] -> Expr) -> m [Expr] -> m Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Expr -> m Expr
go (Expr -> m Expr) -> [Expr] -> m [Expr]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
<$$> [Expr]
ps)              )
    go (POr  [Expr]
ps)       = Expr -> m Expr
f (Expr -> m Expr) -> m Expr -> m Expr
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< ([Expr] -> Expr
POr         ([Expr] -> Expr) -> m [Expr] -> m Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Expr -> m Expr
go (Expr -> m Expr) -> [Expr] -> m [Expr]
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
<$$> [Expr]
ps)              )

mapKVarSubsts :: Visitable t => (KVar -> Subst -> Subst) -> t -> t
mapKVarSubsts :: (KVar -> Subst -> Subst) -> t -> t
mapKVarSubsts KVar -> Subst -> Subst
f          = Visitor () () -> () -> () -> t -> t
forall t a ctx.
(Visitable t, Monoid a) =>
Visitor a ctx -> ctx -> a -> t -> t
trans Visitor () ()
forall p. Visitor () p
kvVis () ()
  where
    kvVis :: Visitor () p
kvVis                = Visitor () p
forall acc ctx. Monoid acc => Visitor acc ctx
defaultVisitor { txExpr :: p -> Expr -> Expr
txExpr = p -> Expr -> Expr
forall p. p -> Expr -> Expr
txK }
    txK :: p -> Expr -> Expr
txK p
_ (PKVar KVar
k Subst
su)   = KVar -> Subst -> Expr
PKVar KVar
k (KVar -> Subst -> Subst
f KVar
k Subst
su)
    txK p
_ (PGrad KVar
k Subst
su GradInfo
i Expr
e) = KVar -> Subst -> GradInfo -> Expr -> Expr
PGrad KVar
k (KVar -> Subst -> Subst
f KVar
k Subst
su) GradInfo
i Expr
e
    txK p
_ Expr
p              = Expr
p

newtype MInt = MInt Integer -- deriving (Eq, NFData)

instance Semigroup MInt where
  MInt SubcId
m <> :: MInt -> MInt -> MInt
<> MInt SubcId
n = SubcId -> MInt
MInt (SubcId
m SubcId -> SubcId -> SubcId
forall a. Num a => a -> a -> a
+ SubcId
n)

instance Monoid MInt where
  mempty :: MInt
mempty  = SubcId -> MInt
MInt SubcId
0
  mappend :: MInt -> MInt -> MInt
mappend = MInt -> MInt -> MInt
forall a. Semigroup a => a -> a -> a
(<>)

size :: Visitable t => t -> Integer
size :: t -> SubcId
size t
t    = SubcId
n
  where
    MInt SubcId
n = Visitor MInt () -> () -> MInt -> t -> MInt
forall t a ctx.
(Visitable t, Monoid a) =>
Visitor a ctx -> ctx -> a -> t -> a
fold Visitor MInt ()
forall ctx. Visitor MInt ctx
szV () MInt
forall a. Monoid a => a
mempty t
t
    szV :: Visitor MInt ctx
szV    = (forall ctx. Visitor MInt ctx
forall acc ctx. Monoid acc => Visitor acc ctx
defaultVisitor :: Visitor MInt t) { accExpr :: ctx -> Expr -> MInt
accExpr = \ ctx
_ Expr
_ -> SubcId -> MInt
MInt SubcId
1 }


lamSize :: Visitable t => t -> Integer
lamSize :: t -> SubcId
lamSize t
t    = SubcId
n
  where
    MInt SubcId
n = Visitor MInt () -> () -> MInt -> t -> MInt
forall t a ctx.
(Visitable t, Monoid a) =>
Visitor a ctx -> ctx -> a -> t -> a
fold Visitor MInt ()
forall ctx. Visitor MInt ctx
szV () MInt
forall a. Monoid a => a
mempty t
t
    szV :: Visitor MInt p
szV    = (forall ctx. Visitor MInt ctx
forall acc ctx. Monoid acc => Visitor acc ctx
defaultVisitor :: Visitor MInt t) { accExpr :: p -> Expr -> MInt
accExpr = p -> Expr -> MInt
forall p. p -> Expr -> MInt
accum }
    accum :: p -> Expr -> MInt
accum p
_ (ELam (Symbol, Sort)
_ Expr
_) = SubcId -> MInt
MInt SubcId
1
    accum p
_ Expr
_          = SubcId -> MInt
MInt SubcId
0

eapps :: Visitable t => t -> [Expr]
eapps :: t -> [Expr]
eapps                 = Visitor [Expr] () -> () -> [Expr] -> t -> [Expr]
forall t a ctx.
(Visitable t, Monoid a) =>
Visitor a ctx -> ctx -> a -> t -> a
fold Visitor [Expr] ()
forall p. Visitor [Expr] p
eappVis () []
  where
    eappVis :: Visitor [Expr] p
eappVis              = (forall t. Visitor [KVar] t
forall acc ctx. Monoid acc => Visitor acc ctx
defaultVisitor :: Visitor [KVar] t) { accExpr :: p -> Expr -> [Expr]
accExpr = p -> Expr -> [Expr]
forall p. p -> Expr -> [Expr]
eapp' }
    eapp' :: p -> Expr -> [Expr]
eapp' p
_ e :: Expr
e@(EApp Expr
_ Expr
_) = [Expr
e]
    eapp' p
_ Expr
_            = []

kvars :: Visitable t => t -> [KVar]
kvars :: t -> [KVar]
kvars                 = Visitor [KVar] () -> () -> [KVar] -> t -> [KVar]
forall t a ctx.
(Visitable t, Monoid a) =>
Visitor a ctx -> ctx -> a -> t -> a
fold Visitor [KVar] ()
forall t. Visitor [KVar] t
kvVis () []
  where
    kvVis :: Visitor [KVar] p
kvVis             = (forall t. Visitor [KVar] t
forall acc ctx. Monoid acc => Visitor acc ctx
defaultVisitor :: Visitor [KVar] t) { accExpr :: p -> Expr -> [KVar]
accExpr = p -> Expr -> [KVar]
forall p. p -> Expr -> [KVar]
kv' }
    kv' :: p -> Expr -> [KVar]
kv' p
_ (PKVar KVar
k Subst
_)     = [KVar
k]
    kv' p
_ (PGrad KVar
k Subst
_ GradInfo
_ Expr
_) = [KVar
k]
    kv' p
_ Expr
_               = []

envKVars :: (TaggedC c a) => BindEnv -> c a -> [KVar]
envKVars :: BindEnv -> c a -> [KVar]
envKVars BindEnv
be c a
c = [[KVar]] -> [KVar]
squish [ SortedReft -> [KVar]
kvs SortedReft
sr |  (Symbol
_, SortedReft
sr) <- BindEnv -> c a -> [(Symbol, SortedReft)]
forall (c :: * -> *) a.
TaggedC c a =>
BindEnv -> c a -> [(Symbol, SortedReft)]
clhs BindEnv
be c a
c]
  where
    squish :: [[KVar]] -> [KVar]
squish    = HashSet KVar -> [KVar]
forall a. HashSet a -> [a]
S.toList  (HashSet KVar -> [KVar])
-> ([[KVar]] -> HashSet KVar) -> [[KVar]] -> [KVar]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [KVar] -> HashSet KVar
forall a. (Eq a, Hashable a) => [a] -> HashSet a
S.fromList ([KVar] -> HashSet KVar)
-> ([[KVar]] -> [KVar]) -> [[KVar]] -> HashSet KVar
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[KVar]] -> [KVar]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
    kvs :: SortedReft -> [KVar]
kvs       = Reft -> [KVar]
forall t. Visitable t => t -> [KVar]
kvars (Reft -> [KVar]) -> (SortedReft -> Reft) -> SortedReft -> [KVar]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SortedReft -> Reft
sr_reft

envKVarsN :: (TaggedC c a) => BindEnv -> c a -> [(KVar, Int)]
envKVarsN :: BindEnv -> c a -> [(KVar, Int)]
envKVarsN BindEnv
be c a
c = [[KVar]] -> [(KVar, Int)]
tally [ SortedReft -> [KVar]
kvs SortedReft
sr |  (Symbol
_, SortedReft
sr) <- BindEnv -> c a -> [(Symbol, SortedReft)]
forall (c :: * -> *) a.
TaggedC c a =>
BindEnv -> c a -> [(Symbol, SortedReft)]
clhs BindEnv
be c a
c]
  where
    tally :: [[KVar]] -> [(KVar, Int)]
tally      = [KVar] -> [(KVar, Int)]
forall k. (Eq k, Hashable k) => [k] -> [(k, Int)]
Misc.count ([KVar] -> [(KVar, Int)])
-> ([[KVar]] -> [KVar]) -> [[KVar]] -> [(KVar, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[KVar]] -> [KVar]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
    kvs :: SortedReft -> [KVar]
kvs        = Reft -> [KVar]
forall t. Visitable t => t -> [KVar]
kvars (Reft -> [KVar]) -> (SortedReft -> Reft) -> SortedReft -> [KVar]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SortedReft -> Reft
sr_reft

rhsKVars :: (TaggedC c a) => c a -> [KVar]
rhsKVars :: c a -> [KVar]
rhsKVars = Expr -> [KVar]
forall t. Visitable t => t -> [KVar]
kvars (Expr -> [KVar]) -> (c a -> Expr) -> c a -> [KVar]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. c a -> Expr
forall (c :: * -> *) a. TaggedC c a => c a -> Expr
crhs -- rhsCs

isKvarC :: (TaggedC c a) => c a -> Bool
isKvarC :: c a -> Bool
isKvarC = (Expr -> Bool) -> [Expr] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Expr -> Bool
isKvar ([Expr] -> Bool) -> (c a -> [Expr]) -> c a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr -> [Expr]
conjuncts (Expr -> [Expr]) -> (c a -> Expr) -> c a -> [Expr]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. c a -> Expr
forall (c :: * -> *) a. TaggedC c a => c a -> Expr
crhs

isConcC :: (TaggedC c a) => c a -> Bool
isConcC :: c a -> Bool
isConcC = (Expr -> Bool) -> [Expr] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all Expr -> Bool
isConc ([Expr] -> Bool) -> (c a -> [Expr]) -> c a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr -> [Expr]
conjuncts (Expr -> [Expr]) -> (c a -> Expr) -> c a -> [Expr]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. c a -> Expr
forall (c :: * -> *) a. TaggedC c a => c a -> Expr
crhs

isKvar :: Expr -> Bool
isKvar :: Expr -> Bool
isKvar (PKVar {}) = Bool
True
isKvar (PGrad {}) = Bool
True
isKvar Expr
_          = Bool
False

isConc :: Expr -> Bool
isConc :: Expr -> Bool
isConc = [KVar] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([KVar] -> Bool) -> (Expr -> [KVar]) -> Expr -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr -> [KVar]
forall t. Visitable t => t -> [KVar]
kvars

stripCasts :: (Visitable t) => t -> t
stripCasts :: t -> t
stripCasts = Visitor () () -> () -> () -> t -> t
forall t a ctx.
(Visitable t, Monoid a) =>
Visitor a ctx -> ctx -> a -> t -> t
trans (Visitor () ()
forall acc ctx. Monoid acc => Visitor acc ctx
defaultVisitor { txExpr :: () -> Expr -> Expr
txExpr = (Expr -> Expr) -> () -> Expr -> Expr
forall a b. a -> b -> a
const Expr -> Expr
go }) () ()
  where
    go :: Expr -> Expr
go (ECst Expr
e Sort
_) = Expr
e
    go Expr
e          = Expr
e

-- stripCasts :: Expr -> Expr
-- stripCasts = mapExpr go
--  where
--    go (ECst e _) = e
--    go e          = e

--------------------------------------------------------------------------------
-- | @CoSub@ is a map from (coercion) ty-vars represented as 'FObj s'
--   to the ty-vars that they should be substituted with. Note the
--   domain and range are both Symbol and not the Int used for real ty-vars.
--------------------------------------------------------------------------------
type CoSub = M.HashMap Symbol Sort 

applyCoSub :: CoSub -> Expr -> Expr
applyCoSub :: CoSub -> Expr -> Expr
applyCoSub CoSub
coSub      = (Expr -> Expr) -> Expr -> Expr
forall t. Visitable t => (Expr -> Expr) -> t -> t
mapExpr Expr -> Expr
fE
  where
    fE :: Expr -> Expr
fE (ECoerc Sort
s Sort
t Expr
e) = Sort -> Sort -> Expr -> Expr
ECoerc  (Sort -> Sort
txS Sort
s) (Sort -> Sort
txS Sort
t) Expr
e
    fE (ELam (Symbol
x,Sort
t) Expr
e) = (Symbol, Sort) -> Expr -> Expr
ELam (Symbol
x, Sort -> Sort
txS Sort
t)         Expr
e
    fE Expr
e              = Expr
e
    txS :: Sort -> Sort
txS               = (Sort -> Sort) -> Sort -> Sort
mapSort Sort -> Sort
fS
    fS :: Sort -> Sort
fS (FObj Symbol
a)       = {- FObj -} (Symbol -> Sort
txV Symbol
a)
    fS Sort
t              = Sort
t
    txV :: Symbol -> Sort
txV Symbol
a             = Sort -> Symbol -> CoSub -> Sort
forall k v. (Eq k, Hashable k) => v -> k -> HashMap k v -> v
M.lookupDefault (Symbol -> Sort
FObj Symbol
a) Symbol
a CoSub
coSub

---------------------------------------------------------------------------------
-- | Visitors over @Sort@
---------------------------------------------------------------------------------
foldSort :: (a -> Sort -> a) -> a -> Sort -> a
foldSort :: (a -> Sort -> a) -> a -> Sort -> a
foldSort a -> Sort -> a
f = a -> Sort -> a
step
  where
    step :: a -> Sort -> a
step a
b Sort
t           = a -> Sort -> a
go (a -> Sort -> a
f a
b Sort
t) Sort
t
    go :: a -> Sort -> a
go a
b (FFunc Sort
t1 Sort
t2) = (a -> Sort -> a) -> a -> [Sort] -> a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
L.foldl' a -> Sort -> a
step a
b [Sort
t1, Sort
t2]
    go a
b (FApp Sort
t1 Sort
t2)  = (a -> Sort -> a) -> a -> [Sort] -> a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
L.foldl' a -> Sort -> a
step a
b [Sort
t1, Sort
t2]
    go a
b (FAbs Int
_ Sort
t)    = a -> Sort -> a
go a
b Sort
t
    go a
b Sort
_             = a
b

mapSort :: (Sort -> Sort) -> Sort -> Sort
mapSort :: (Sort -> Sort) -> Sort -> Sort
mapSort Sort -> Sort
f = Sort -> Sort
step
  where
    step :: Sort -> Sort
step !Sort
x           = Sort -> Sort
go (Sort -> Sort
f Sort
x)
    go :: Sort -> Sort
go !(FFunc Sort
t1 Sort
t2) = Sort -> Sort -> Sort
FFunc (Sort -> Sort
step Sort
t1) (Sort -> Sort
step Sort
t2)
    go !(FApp Sort
t1 Sort
t2)  = Sort -> Sort -> Sort
FApp  (Sort -> Sort
step Sort
t1) (Sort -> Sort
step Sort
t2)
    go !(FAbs Int
i Sort
t)    = Int -> Sort -> Sort
FAbs Int
i (Sort -> Sort
step Sort
t)
    go !Sort
t             = Sort
t

foldDataDecl :: (a -> Sort -> a) -> a -> DataDecl -> a
foldDataDecl :: (a -> Sort -> a) -> a -> DataDecl -> a
foldDataDecl a -> Sort -> a
f a
acc = (a -> Sort -> a) -> a -> [Sort] -> a
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
L.foldl' a -> Sort -> a
f a
acc ([Sort] -> a) -> (DataDecl -> [Sort]) -> DataDecl -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DataDecl -> [Sort]
dataDeclSorts

dataDeclSorts :: DataDecl -> [Sort]
dataDeclSorts :: DataDecl -> [Sort]
dataDeclSorts = (DataCtor -> [Sort]) -> [DataCtor] -> [Sort]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap DataCtor -> [Sort]
dataCtorSorts ([DataCtor] -> [Sort])
-> (DataDecl -> [DataCtor]) -> DataDecl -> [Sort]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DataDecl -> [DataCtor]
ddCtors

dataCtorSorts :: DataCtor -> [Sort]
dataCtorSorts :: DataCtor -> [Sort]
dataCtorSorts = (DataField -> Sort) -> [DataField] -> [Sort]
forall a b. (a -> b) -> [a] -> [b]
map DataField -> Sort
dfSort ([DataField] -> [Sort])
-> (DataCtor -> [DataField]) -> DataCtor -> [Sort]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DataCtor -> [DataField]
dcFields
---------------------------------------------------------------
-- | String Constants -----------------------------------------
---------------------------------------------------------------

-- symConstLits    :: FInfo a -> [(Symbol, Sort)]
-- symConstLits fi = [(symbol c, strSort) | c <- symConsts fi]

class SymConsts a where
  symConsts :: a -> [SymConst]

-- instance  SymConsts (FInfo a) where
instance (SymConsts (c a)) => SymConsts (GInfo c a) where
  symConsts :: GInfo c a -> [SymConst]
symConsts GInfo c a
fi = [SymConst] -> [SymConst]
forall a. Ord a => [a] -> [a]
Misc.sortNub ([SymConst] -> [SymConst]) -> [SymConst] -> [SymConst]
forall a b. (a -> b) -> a -> b
$ [SymConst]
csLits [SymConst] -> [SymConst] -> [SymConst]
forall a. [a] -> [a] -> [a]
++ [SymConst]
bsLits [SymConst] -> [SymConst] -> [SymConst]
forall a. [a] -> [a] -> [a]
++ [SymConst]
qsLits
    where
      csLits :: [SymConst]
csLits   = (c a -> [SymConst]) -> [c a] -> [SymConst]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap c a -> [SymConst]
forall a. SymConsts a => a -> [SymConst]
symConsts ([c a] -> [SymConst]) -> [c a] -> [SymConst]
forall a b. (a -> b) -> a -> b
$ HashMap SubcId (c a) -> [c a]
forall k v. HashMap k v -> [v]
M.elems  (HashMap SubcId (c a) -> [c a]) -> HashMap SubcId (c a) -> [c a]
forall a b. (a -> b) -> a -> b
$  GInfo c a -> HashMap SubcId (c a)
forall (c :: * -> *) a. GInfo c a -> HashMap SubcId (c a)
cm    GInfo c a
fi
      bsLits :: [SymConst]
bsLits   = BindEnv -> [SymConst]
forall a. SymConsts a => a -> [SymConst]
symConsts           (BindEnv -> [SymConst]) -> BindEnv -> [SymConst]
forall a b. (a -> b) -> a -> b
$ GInfo c a -> BindEnv
forall (c :: * -> *) a. GInfo c a -> BindEnv
bs                GInfo c a
fi
      qsLits :: [SymConst]
qsLits   = (Expr -> [SymConst]) -> [Expr] -> [SymConst]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Expr -> [SymConst]
forall a. SymConsts a => a -> [SymConst]
symConsts ([Expr] -> [SymConst]) -> [Expr] -> [SymConst]
forall a b. (a -> b) -> a -> b
$ Qualifier -> Expr
qBody   (Qualifier -> Expr) -> [Qualifier] -> [Expr]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> GInfo c a -> [Qualifier]
forall (c :: * -> *) a. GInfo c a -> [Qualifier]
quals GInfo c a
fi

instance SymConsts BindEnv where
  symConsts :: BindEnv -> [SymConst]
symConsts    = ((Symbol, SortedReft) -> [SymConst])
-> [(Symbol, SortedReft)] -> [SymConst]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (SortedReft -> [SymConst]
forall a. SymConsts a => a -> [SymConst]
symConsts (SortedReft -> [SymConst])
-> ((Symbol, SortedReft) -> SortedReft)
-> (Symbol, SortedReft)
-> [SymConst]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Symbol, SortedReft) -> SortedReft
forall a b. (a, b) -> b
snd) ([(Symbol, SortedReft)] -> [SymConst])
-> (BindEnv -> [(Symbol, SortedReft)]) -> BindEnv -> [SymConst]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HashMap Int (Symbol, SortedReft) -> [(Symbol, SortedReft)]
forall k v. HashMap k v -> [v]
M.elems (HashMap Int (Symbol, SortedReft) -> [(Symbol, SortedReft)])
-> (BindEnv -> HashMap Int (Symbol, SortedReft))
-> BindEnv
-> [(Symbol, SortedReft)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BindEnv -> HashMap Int (Symbol, SortedReft)
forall a. SizedEnv a -> BindMap a
beBinds

instance SymConsts (SubC a) where
  symConsts :: SubC a -> [SymConst]
symConsts SubC a
c  = SortedReft -> [SymConst]
forall a. SymConsts a => a -> [SymConst]
symConsts (SubC a -> SortedReft
forall a. SubC a -> SortedReft
slhs SubC a
c) [SymConst] -> [SymConst] -> [SymConst]
forall a. [a] -> [a] -> [a]
++
                 SortedReft -> [SymConst]
forall a. SymConsts a => a -> [SymConst]
symConsts (SubC a -> SortedReft
forall a. SubC a -> SortedReft
srhs SubC a
c)

instance SymConsts (SimpC a) where
  symConsts :: SimpC a -> [SymConst]
symConsts SimpC a
c  = Expr -> [SymConst]
forall a. SymConsts a => a -> [SymConst]
symConsts (SimpC a -> Expr
forall (c :: * -> *) a. TaggedC c a => c a -> Expr
crhs SimpC a
c)

instance SymConsts SortedReft where
  symConsts :: SortedReft -> [SymConst]
symConsts = Reft -> [SymConst]
forall a. SymConsts a => a -> [SymConst]
symConsts (Reft -> [SymConst])
-> (SortedReft -> Reft) -> SortedReft -> [SymConst]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SortedReft -> Reft
sr_reft

instance SymConsts Reft where
  symConsts :: Reft -> [SymConst]
symConsts (Reft (Symbol
_, Expr
ra)) = Expr -> [SymConst]
forall t. Visitable t => t -> [SymConst]
getSymConsts Expr
ra


instance SymConsts Expr where
  symConsts :: Expr -> [SymConst]
symConsts = Expr -> [SymConst]
forall t. Visitable t => t -> [SymConst]
getSymConsts

getSymConsts :: Visitable t => t -> [SymConst]
getSymConsts :: t -> [SymConst]
getSymConsts         = Visitor [SymConst] () -> () -> [SymConst] -> t -> [SymConst]
forall t a ctx.
(Visitable t, Monoid a) =>
Visitor a ctx -> ctx -> a -> t -> a
fold Visitor [SymConst] ()
forall p. Visitor [SymConst] p
scVis () []
  where
    scVis :: Visitor [SymConst] p
scVis            = (forall p. Visitor [SymConst] p
forall acc ctx. Monoid acc => Visitor acc ctx
defaultVisitor :: Visitor [SymConst] t)  { accExpr :: p -> Expr -> [SymConst]
accExpr = p -> Expr -> [SymConst]
forall p. p -> Expr -> [SymConst]
sc }
    sc :: p -> Expr -> [SymConst]
sc p
_ (ESym SymConst
c)    = [SymConst
c]
    sc p
_ Expr
_           = []