-- | This module has various utility functions for accessing queries.
--   TODO: move the "clients" in Visitors into this module.

module Language.Fixpoint.Types.Utils (
  -- * Domain of a kvar
    kvarDomain

  -- * Free variables in a refinement
  , reftFreeVars

  -- * Deconstruct a SortedReft
  , sortedReftConcKVars
  ) where

import qualified Data.HashMap.Strict                  as M
import qualified Data.HashSet                         as S

import           Language.Fixpoint.Misc
import           Language.Fixpoint.Types.Names
import           Language.Fixpoint.Types.Refinements
import           Language.Fixpoint.Types.Environments
import           Language.Fixpoint.Types.Constraints

--------------------------------------------------------------------------------
-- | Compute the domain of a kvar
--------------------------------------------------------------------------------
kvarDomain :: SInfo a -> KVar -> [Symbol]
--------------------------------------------------------------------------------
kvarDomain :: SInfo a -> KVar -> [Symbol]
kvarDomain SInfo a
si KVar
k = BindEnv -> WfC a -> [Symbol]
forall a. BindEnv -> WfC a -> [Symbol]
domain (SInfo a -> BindEnv
forall (c :: * -> *) a. GInfo c a -> BindEnv
bs SInfo a
si) (SInfo a -> KVar -> WfC a
forall a. SInfo a -> KVar -> WfC a
getWfC SInfo a
si KVar
k)

domain :: BindEnv -> WfC a -> [Symbol]
domain :: BindEnv -> WfC a -> [Symbol]
domain BindEnv
be WfC a
wfc = (Symbol, Sort, KVar) -> Symbol
forall a b c. (a, b, c) -> a
fst3 (WfC a -> (Symbol, Sort, KVar)
forall a. WfC a -> (Symbol, Sort, KVar)
wrft WfC a
wfc) Symbol -> [Symbol] -> [Symbol]
forall a. a -> [a] -> [a]
: ((Symbol, SortedReft) -> Symbol)
-> [(Symbol, SortedReft)] -> [Symbol]
forall a b. (a -> b) -> [a] -> [b]
map (Symbol, SortedReft) -> Symbol
forall a b. (a, b) -> a
fst (BindEnv -> IBindEnv -> [(Symbol, SortedReft)]
envCs BindEnv
be (IBindEnv -> [(Symbol, SortedReft)])
-> IBindEnv -> [(Symbol, SortedReft)]
forall a b. (a -> b) -> a -> b
$ WfC a -> IBindEnv
forall a. WfC a -> IBindEnv
wenv WfC a
wfc)

getWfC :: SInfo a -> KVar -> WfC a
getWfC :: SInfo a -> KVar -> WfC a
getWfC SInfo a
si KVar
k = SInfo a -> HashMap KVar (WfC a)
forall (c :: * -> *) a. GInfo c a -> HashMap KVar (WfC a)
ws SInfo a
si HashMap KVar (WfC a) -> KVar -> WfC a
forall k v.
(Eq k, Hashable k, HasCallStack) =>
HashMap k v -> k -> v
M.! KVar
k

--------------------------------------------------------------------------------
-- | Free variables of a refinement
--------------------------------------------------------------------------------
--TODO deduplicate (also in Solver/UniqifyBinds)
reftFreeVars :: Reft -> S.HashSet Symbol
reftFreeVars :: Reft -> HashSet Symbol
reftFreeVars r :: Reft
r@(Reft (Symbol
v, Expr
_)) = Symbol -> HashSet Symbol -> HashSet Symbol
forall a. (Eq a, Hashable a) => a -> HashSet a -> HashSet a
S.delete Symbol
v (HashSet Symbol -> HashSet Symbol)
-> HashSet Symbol -> HashSet Symbol
forall a b. (a -> b) -> a -> b
$ [Symbol] -> HashSet Symbol
forall a. (Eq a, Hashable a) => [a] -> HashSet a
S.fromList ([Symbol] -> HashSet Symbol) -> [Symbol] -> HashSet Symbol
forall a b. (a -> b) -> a -> b
$ Reft -> [Symbol]
forall a. Subable a => a -> [Symbol]
syms Reft
r

--------------------------------------------------------------------------------
-- | Split a SortedReft into its concrete and KVar components
--------------------------------------------------------------------------------
sortedReftConcKVars :: Symbol -> SortedReft -> ([Pred], [KVSub], [KVSub])
sortedReftConcKVars :: Symbol -> SortedReft -> ([Expr], [KVSub], [KVSub])
sortedReftConcKVars Symbol
x SortedReft
sr = [Expr]
-> [KVSub]
-> [KVSub]
-> [(Symbol, Expr)]
-> ([Expr], [KVSub], [KVSub])
go [] [] [] [(Symbol, Expr)]
ves
  where
    ves :: [(Symbol, Expr)]
ves                  = [(Symbol
v, Expr
p Expr -> (Symbol, Expr) -> Expr
forall a. Subable a => a -> (Symbol, Expr) -> a
`subst1` (Symbol
v, Symbol -> Expr
forall a. Symbolic a => a -> Expr
eVar Symbol
x)) | Reft (Symbol
v, Expr
p) <- [Reft]
rs ]
    rs :: [Reft]
rs                   = Reft -> [Reft]
reftConjuncts (SortedReft -> Reft
sr_reft SortedReft
sr)
    t :: Sort
t                    = SortedReft -> Sort
sr_sort SortedReft
sr

    go :: [Expr]
-> [KVSub]
-> [KVSub]
-> [(Symbol, Expr)]
-> ([Expr], [KVSub], [KVSub])
go [Expr]
ps [KVSub]
ks [KVSub]
gs ((Symbol
v, PKVar KVar
k Subst
su    ):[(Symbol, Expr)]
xs) = [Expr]
-> [KVSub]
-> [KVSub]
-> [(Symbol, Expr)]
-> ([Expr], [KVSub], [KVSub])
go [Expr]
ps (Symbol -> Sort -> KVar -> Subst -> KVSub
KVS Symbol
v Sort
t KVar
k Subst
suKVSub -> [KVSub] -> [KVSub]
forall a. a -> [a] -> [a]
:[KVSub]
ks) [KVSub]
gs [(Symbol, Expr)]
xs 
    go [Expr]
ps [KVSub]
ks [KVSub]
gs ((Symbol
v, PGrad KVar
k Subst
su GradInfo
_ Expr
_):[(Symbol, Expr)]
xs) = [Expr]
-> [KVSub]
-> [KVSub]
-> [(Symbol, Expr)]
-> ([Expr], [KVSub], [KVSub])
go [Expr]
ps [KVSub]
ks (Symbol -> Sort -> KVar -> Subst -> KVSub
KVS Symbol
v Sort
t KVar
k Subst
suKVSub -> [KVSub] -> [KVSub]
forall a. a -> [a] -> [a]
:[KVSub]
gs) [(Symbol, Expr)]
xs 
    go [Expr]
ps [KVSub]
ks [KVSub]
gs ((Symbol
_, Expr
p):[(Symbol, Expr)]
xs)              = [Expr]
-> [KVSub]
-> [KVSub]
-> [(Symbol, Expr)]
-> ([Expr], [KVSub], [KVSub])
go (Expr
pExpr -> [Expr] -> [Expr]
forall a. a -> [a] -> [a]
:[Expr]
ps) [KVSub]
ks [KVSub]
gs [(Symbol, Expr)]
xs 
    go [Expr]
ps [KVSub]
ks [KVSub]
gs []                       = ([Expr]
ps, [KVSub]
ks, [KVSub]
gs)