-- | 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

  -- * Operators on DataDecl
  , checkRegular
  , orderDeclarations

  ) 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
import           Language.Fixpoint.Types.Sorts
import qualified Language.Fixpoint.Misc as Misc

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

domain :: BindEnv a -> WfC a -> [Symbol]
domain :: forall a. BindEnv a -> WfC a -> [Symbol]
domain BindEnv a
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 a -> IBindEnv -> [(Symbol, SortedReft)]
forall a. BindEnv a -> IBindEnv -> [(Symbol, SortedReft)]
envCs BindEnv a
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 :: forall a. 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, ExprV Symbol
_)) = 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 -> ([ExprV Symbol], [KVSub], [KVSub])
sortedReftConcKVars Symbol
x SortedReft
sr = [ExprV Symbol]
-> [KVSub]
-> [KVSub]
-> [(Symbol, ExprV Symbol)]
-> ([ExprV Symbol], [KVSub], [KVSub])
go [] [] [] [(Symbol, ExprV Symbol)]
ves
  where
    ves :: [(Symbol, ExprV Symbol)]
ves                  = [(Symbol
v, ExprV Symbol
p ExprV Symbol -> (Symbol, ExprV Symbol) -> ExprV Symbol
forall a. Subable a => a -> (Symbol, ExprV Symbol) -> a
`subst1` (Symbol
v, Symbol -> ExprV Symbol
forall a. Symbolic a => a -> ExprV Symbol
eVar Symbol
x)) | Reft (Symbol
v, ExprV Symbol
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 :: [ExprV Symbol]
-> [KVSub]
-> [KVSub]
-> [(Symbol, ExprV Symbol)]
-> ([ExprV Symbol], [KVSub], [KVSub])
go [ExprV Symbol]
ps [KVSub]
ks [KVSub]
gs ((Symbol
v, PKVar KVar
k SubstV Symbol
su    ):[(Symbol, ExprV Symbol)]
xs) = [ExprV Symbol]
-> [KVSub]
-> [KVSub]
-> [(Symbol, ExprV Symbol)]
-> ([ExprV Symbol], [KVSub], [KVSub])
go [ExprV Symbol]
ps (Symbol -> Sort -> KVar -> SubstV Symbol -> KVSub
KVS Symbol
v Sort
t KVar
k SubstV Symbol
suKVSub -> [KVSub] -> [KVSub]
forall a. a -> [a] -> [a]
:[KVSub]
ks) [KVSub]
gs [(Symbol, ExprV Symbol)]
xs
    go [ExprV Symbol]
ps [KVSub]
ks [KVSub]
gs ((Symbol
v, PGrad KVar
k SubstV Symbol
su GradInfo
_ ExprV Symbol
_):[(Symbol, ExprV Symbol)]
xs) = [ExprV Symbol]
-> [KVSub]
-> [KVSub]
-> [(Symbol, ExprV Symbol)]
-> ([ExprV Symbol], [KVSub], [KVSub])
go [ExprV Symbol]
ps [KVSub]
ks (Symbol -> Sort -> KVar -> SubstV Symbol -> KVSub
KVS Symbol
v Sort
t KVar
k SubstV Symbol
suKVSub -> [KVSub] -> [KVSub]
forall a. a -> [a] -> [a]
:[KVSub]
gs) [(Symbol, ExprV Symbol)]
xs
    go [ExprV Symbol]
ps [KVSub]
ks [KVSub]
gs ((Symbol
_, ExprV Symbol
p):[(Symbol, ExprV Symbol)]
xs)              = [ExprV Symbol]
-> [KVSub]
-> [KVSub]
-> [(Symbol, ExprV Symbol)]
-> ([ExprV Symbol], [KVSub], [KVSub])
go (ExprV Symbol
pExprV Symbol -> [ExprV Symbol] -> [ExprV Symbol]
forall a. a -> [a] -> [a]
:[ExprV Symbol]
ps) [KVSub]
ks [KVSub]
gs [(Symbol, ExprV Symbol)]
xs
    go [ExprV Symbol]
ps [KVSub]
ks [KVSub]
gs []                       = ([ExprV Symbol]
ps, [KVSub]
ks, [KVSub]
gs)


-------------------------------------------------------------------------------
-- | @checkRegular ds@ returns the subset of ds that are _not_ regular
-------------------------------------------------------------------------------
checkRegular :: [DataDecl] -> [DataDecl]
-------------------------------------------------------------------------------
checkRegular :: [DataDecl] -> [DataDecl]
checkRegular [DataDecl]
ds = [[DataDecl]] -> [DataDecl]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [ [DataDecl]
ds' | [DataDecl]
ds' <- [DataDecl] -> [[DataDecl]]
orderDeclarations [DataDecl]
ds, Bool -> Bool
not ([DataDecl] -> Bool
isRegular [DataDecl]
ds')]

-------------------------------------------------------------------------------
-- | @isRegular [d1,...]@ gets a non-empty list of mut-recursive datadecls
-------------------------------------------------------------------------------
isRegular :: [DataDecl] -> Bool
-------------------------------------------------------------------------------

isRegular :: [DataDecl] -> Bool
isRegular []       = [Char] -> Bool
forall a. HasCallStack => [Char] -> a
error [Char]
"impossible: isRegular"
isRegular ds :: [DataDecl]
ds@(DataDecl
d:[DataDecl]
_) = (DataDecl -> Bool) -> [DataDecl] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\DataDecl
d' -> DataDecl -> Int
ddVars DataDecl
d' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
nArgs) [DataDecl]
ds   -- same number of tyArgs
                  Bool -> Bool -> Bool
&& ((Sort, [Sort]) -> Bool) -> [(Sort, [Sort])] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Sort, [Sort]) -> Bool
isRegApp [(Sort, [Sort])]
fldSortApps         -- 'regular' application (tc @0 ... @n)
  where
    nArgs :: Int
nArgs          = DataDecl -> Int
ddVars DataDecl
d
    tcs :: HashSet Symbol
tcs            = [Symbol] -> HashSet Symbol
forall a. (Eq a, Hashable a) => [a] -> HashSet a
S.fromList ( FTycon -> Symbol
forall a. Symbolic a => a -> Symbol
symbol (FTycon -> Symbol) -> (DataDecl -> FTycon) -> DataDecl -> Symbol
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DataDecl -> FTycon
ddTyCon (DataDecl -> Symbol) -> [DataDecl] -> [Symbol]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [DataDecl]
ds)
    fldSortApps :: [(Sort, [Sort])]
fldSortApps    = [ (Sort
c,[Sort]
ts) | DataDecl
d'          <- [DataDecl]
ds
                              , DataCtor
ctor        <- DataDecl -> [DataCtor]
ddCtors DataDecl
d'
                              , DField LocSymbol
_ Sort
t  <- DataCtor -> [DataField]
dcFields DataCtor
ctor
                              , (Sort
c, [Sort]
ts)     <- Sort -> [(Sort, [Sort])]
sortApps Sort
t
                     ]
    isRegApp :: (Sort, [Sort]) -> Bool
isRegApp (Sort, [Sort])
cts   = case (Sort, [Sort])
cts of
                        (FTC FTycon
c, [Sort]
ts) -> Bool -> Bool
not (Symbol -> HashSet Symbol -> Bool
forall a. (Eq a, Hashable a) => a -> HashSet a -> Bool
S.member (FTycon -> Symbol
forall a. Symbolic a => a -> Symbol
symbol FTycon
c) HashSet Symbol
tcs) Bool -> Bool -> Bool
|| Int -> [Sort] -> Bool
isRegularArgs Int
nArgs [Sort]
ts
                        (Sort, [Sort])
_           -> Bool
False

isRegularArgs :: Int -> [Sort] -> Bool
isRegularArgs :: Int -> [Sort] -> Bool
isRegularArgs Int
n [Sort]
ts = [Sort]
ts [Sort] -> [Sort] -> Bool
forall a. Eq a => a -> a -> Bool
== [Int -> Sort
FVar Int
i | Int
i <- [Int
0 .. (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)]]

type SortApp = (Sort, [Sort])

sortApps :: Sort -> [SortApp]
sortApps :: Sort -> [(Sort, [Sort])]
sortApps = Sort -> [(Sort, [Sort])]
go
  where
    go :: Sort -> [(Sort, [Sort])]
go t :: Sort
t@FApp {}     = (Sort
f, [Sort]
ts) (Sort, [Sort]) -> [(Sort, [Sort])] -> [(Sort, [Sort])]
forall a. a -> [a] -> [a]
: (Sort -> [(Sort, [Sort])]) -> [Sort] -> [(Sort, [Sort])]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Sort -> [(Sort, [Sort])]
go [Sort]
ts where (Sort
f, [Sort]
ts) = Sort -> (Sort, [Sort])
splitApp Sort
t
    go (FFunc Sort
t1 Sort
t2) = Sort -> [(Sort, [Sort])]
go Sort
t1 [(Sort, [Sort])] -> [(Sort, [Sort])] -> [(Sort, [Sort])]
forall a. [a] -> [a] -> [a]
++ Sort -> [(Sort, [Sort])]
go Sort
t2
    go (FAbs Int
_ Sort
t)    = Sort -> [(Sort, [Sort])]
go Sort
t
    go Sort
_             = []

splitApp :: Sort -> SortApp
splitApp :: Sort -> (Sort, [Sort])
splitApp = [Sort] -> Sort -> (Sort, [Sort])
go []
  where
    go :: [Sort] -> Sort -> (Sort, [Sort])
go [Sort]
stk (FApp Sort
t1 Sort
t2) = [Sort] -> Sort -> (Sort, [Sort])
go (Sort
t2Sort -> [Sort] -> [Sort]
forall a. a -> [a] -> [a]
:[Sort]
stk) Sort
t1
    go [Sort]
stk Sort
t            = (Sort
t, [Sort]
stk)

--------------------------------------------------------------------------------
-- | 'orderDeclarations' sorts the data declarations such that each declarations
--   only refers to preceding ones.
--------------------------------------------------------------------------------
orderDeclarations :: [DataDecl] -> [[DataDecl]]
--------------------------------------------------------------------------------
orderDeclarations :: [DataDecl] -> [[DataDecl]]
orderDeclarations [DataDecl]
ds = {- reverse -} (DataDecl -> (FTycon, [FTycon])) -> [DataDecl] -> [[DataDecl]]
forall v a. Ord v => (a -> (v, [v])) -> [a] -> [[a]]
Misc.sccsWith DataDecl -> (FTycon, [FTycon])
f [DataDecl]
ds
  where
    dM :: HashMap FTycon DataDecl
dM               = [(FTycon, DataDecl)] -> HashMap FTycon DataDecl
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
M.fromList [(DataDecl -> FTycon
ddTyCon DataDecl
d, DataDecl
d) | DataDecl
d <- [DataDecl]
ds]
    f :: DataDecl -> (FTycon, [FTycon])
f DataDecl
d              = (DataDecl -> FTycon
ddTyCon DataDecl
d, HashMap FTycon DataDecl -> DataDecl -> [FTycon]
dataDeclDeps HashMap FTycon DataDecl
dM DataDecl
d)

dataDeclDeps :: M.HashMap FTycon DataDecl -> DataDecl -> [FTycon]
dataDeclDeps :: HashMap FTycon DataDecl -> DataDecl -> [FTycon]
dataDeclDeps HashMap FTycon DataDecl
dM = (FTycon -> Bool) -> [FTycon] -> [FTycon]
forall a. (a -> Bool) -> [a] -> [a]
filter (FTycon -> HashMap FTycon DataDecl -> Bool
forall k a. (Eq k, Hashable k) => k -> HashMap k a -> Bool
`M.member` HashMap FTycon DataDecl
dM) ([FTycon] -> [FTycon])
-> (DataDecl -> [FTycon]) -> DataDecl -> [FTycon]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [FTycon] -> [FTycon]
forall a. Ord a => [a] -> [a]
Misc.sortNub ([FTycon] -> [FTycon])
-> (DataDecl -> [FTycon]) -> DataDecl -> [FTycon]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DataDecl -> [FTycon]
dataDeclFTycons

dataDeclFTycons :: DataDecl -> [FTycon]
dataDeclFTycons :: DataDecl -> [FTycon]
dataDeclFTycons = (DataCtor -> [FTycon]) -> [DataCtor] -> [FTycon]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap DataCtor -> [FTycon]
dataCtorFTycons ([DataCtor] -> [FTycon])
-> (DataDecl -> [DataCtor]) -> DataDecl -> [FTycon]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DataDecl -> [DataCtor]
ddCtors

dataCtorFTycons :: DataCtor -> [FTycon]
dataCtorFTycons :: DataCtor -> [FTycon]
dataCtorFTycons = (DataField -> [FTycon]) -> [DataField] -> [FTycon]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap DataField -> [FTycon]
dataFieldFTycons ([DataField] -> [FTycon])
-> (DataCtor -> [DataField]) -> DataCtor -> [FTycon]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DataCtor -> [DataField]
dcFields

dataFieldFTycons :: DataField -> [FTycon]
dataFieldFTycons :: DataField -> [FTycon]
dataFieldFTycons = Sort -> [FTycon]
sortFTycons (Sort -> [FTycon]) -> (DataField -> Sort) -> DataField -> [FTycon]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DataField -> Sort
dfSort

sortFTycons :: Sort -> [FTycon]
sortFTycons :: Sort -> [FTycon]
sortFTycons = Sort -> [FTycon]
go
  where
    go :: Sort -> [FTycon]
go (FTC FTycon
c)       = [FTycon
c]
    go (FApp  Sort
t1 Sort
t2) = Sort -> [FTycon]
go Sort
t1 [FTycon] -> [FTycon] -> [FTycon]
forall a. [a] -> [a] -> [a]
++ Sort -> [FTycon]
go Sort
t2
    go (FFunc Sort
t1 Sort
t2) = Sort -> [FTycon]
go Sort
t1 [FTycon] -> [FTycon] -> [FTycon]
forall a. [a] -> [a] -> [a]
++ Sort -> [FTycon]
go Sort
t2
    go (FAbs Int
_ Sort
t)    = Sort -> [FTycon]
go Sort
t
    go Sort
_             = []