{-# LANGUAGE CPP                        #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE NoMonomorphismRestriction  #-}
{-# LANGUAGE OverloadedStrings          #-}
{-# LANGUAGE UndecidableInstances       #-}
{-# LANGUAGE MultiParamTypeClasses      #-}
{-# LANGUAGE TypeOperators              #-}
{-# LANGUAGE GADTs                      #-}
{-# LANGUAGE PatternGuards              #-}
{-# LANGUAGE DeriveGeneric              #-}
{-# LANGUAGE DeriveDataTypeable         #-}
{-# LANGUAGE TupleSections              #-}

-- | This module contains the top-level SOLUTION data types,
--   including various indices used for solving.

module Language.Fixpoint.Types.Graduals (
  uniquify,

  makeSolutions,

  GSol,

  Gradual (..)
  ) where

import Language.Fixpoint.Types.Refinements
import Language.Fixpoint.Types.Constraints
import Language.Fixpoint.Types.Config
import Language.Fixpoint.Types.PrettyPrint
import Language.Fixpoint.Types.Environments
import Language.Fixpoint.Types.Substitutions
import Language.Fixpoint.Types.Visitor
import Language.Fixpoint.Types.Spans
import Language.Fixpoint.Types.Theories
import Language.Fixpoint.Types.Names        (gradIntSymbol, tidySymbol)
import Language.Fixpoint.Misc               (allCombinations, errorstar)

import Control.DeepSeq

import qualified Data.HashMap.Strict       as M
import qualified Data.List                 as L

import Control.Monad.State.Lazy
import Data.Maybe (fromMaybe)
#if !MIN_VERSION_base(4,14,0)
import Data.Semigroup (Semigroup (..))
#endif

import qualified Language.Fixpoint.SortCheck       as So
import Language.Fixpoint.Solver.Sanitize (symbolEnv)


data GSol = GSol !SymEnv !(M.HashMap KVar (Expr, GradInfo))

instance Semigroup GSol where
  (GSol SymEnv
e1 HashMap KVar (Expr, GradInfo)
m1) <> :: GSol -> GSol -> GSol
<> (GSol SymEnv
e2 HashMap KVar (Expr, GradInfo)
m2) = SymEnv -> HashMap KVar (Expr, GradInfo) -> GSol
GSol (SymEnv
e1 SymEnv -> SymEnv -> SymEnv
forall a. Semigroup a => a -> a -> a
<> SymEnv
e2) (HashMap KVar (Expr, GradInfo)
m1 HashMap KVar (Expr, GradInfo)
-> HashMap KVar (Expr, GradInfo) -> HashMap KVar (Expr, GradInfo)
forall a. Semigroup a => a -> a -> a
<> HashMap KVar (Expr, GradInfo)
m2)

instance Monoid GSol where
  mempty :: GSol
mempty = SymEnv -> HashMap KVar (Expr, GradInfo) -> GSol
GSol SymEnv
forall a. Monoid a => a
mempty HashMap KVar (Expr, GradInfo)
forall a. Monoid a => a
mempty

instance Show GSol where
  show :: GSol -> String
show (GSol SymEnv
_ HashMap KVar (Expr, GradInfo)
m) = String
"GSOL = \n" String -> ShowS
forall a. [a] -> [a] -> [a]
++ [String] -> String
unlines ((\(KVar
k,(Expr
e, GradInfo
i)) -> KVar -> String
forall a. PPrint a => a -> String
showpp KVar
k String -> ShowS
forall a. [a] -> [a] -> [a]
++ GradInfo -> String
forall a. Show a => a -> String
showInfo GradInfo
i String -> ShowS
forall a. [a] -> [a] -> [a]
++  String
" |-> " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Expr -> String
forall a. PPrint a => a -> String
showpp (Expr -> Expr
forall a. Subable a => a -> a
tx Expr
e)) ((KVar, (Expr, GradInfo)) -> String)
-> [(KVar, (Expr, GradInfo))] -> [String]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> HashMap KVar (Expr, GradInfo) -> [(KVar, (Expr, GradInfo))]
forall k v. HashMap k v -> [(k, v)]
M.toList HashMap KVar (Expr, GradInfo)
m)
    where
      tx :: a -> a
tx a
e = Subst -> a -> a
forall a. Subable a => Subst -> a -> a
subst ([(Symbol, Expr)] -> Subst
mkSubst ([(Symbol, Expr)] -> Subst) -> [(Symbol, Expr)] -> Subst
forall a b. (a -> b) -> a -> b
$ [(Symbol
x, Symbol -> Expr
EVar (Symbol -> Expr) -> Symbol -> Expr
forall a b. (a -> b) -> a -> b
$ Symbol -> Symbol
tidySymbol Symbol
x) | Symbol
x <- a -> [Symbol]
forall a. Subable a => a -> [Symbol]
syms a
e]) a
e
      showInfo :: a -> String
showInfo a
i = a -> String
forall a. Show a => a -> String
show a
i


makeSolutions :: (NFData a, Fixpoint a, Show a)
              => Config -> SInfo a
              -> [(KVar, (GWInfo, [[Expr]]))]
              -> Maybe [GSol]

makeSolutions :: Config -> SInfo a -> [(KVar, (GWInfo, [[Expr]]))] -> Maybe [GSol]
makeSolutions Config
_ SInfo a
_ []
  = Maybe [GSol]
forall a. Maybe a
Nothing
makeSolutions Config
cfg SInfo a
fi [(KVar, (GWInfo, [[Expr]]))]
kes
  = [GSol] -> Maybe [GSol]
forall a. a -> Maybe a
Just ([GSol] -> Maybe [GSol]) -> [GSol] -> Maybe [GSol]
forall a b. (a -> b) -> a -> b
$ ([(KVar, (Expr, GradInfo))] -> GSol)
-> [[(KVar, (Expr, GradInfo))]] -> [GSol]
forall a b. (a -> b) -> [a] -> [b]
map (SymEnv -> HashMap KVar (Expr, GradInfo) -> GSol
GSol SymEnv
env (HashMap KVar (Expr, GradInfo) -> GSol)
-> ([(KVar, (Expr, GradInfo))] -> HashMap KVar (Expr, GradInfo))
-> [(KVar, (Expr, GradInfo))]
-> GSol
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(KVar, (Expr, GradInfo))] -> HashMap KVar (Expr, GradInfo)
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
M.fromList) ([[(KVar, (Expr, GradInfo))]] -> [[(KVar, (Expr, GradInfo))]]
forall a. [[a]] -> [[a]]
allCombinations ((KVar, (GWInfo, [[Expr]])) -> [(KVar, (Expr, GradInfo))]
forall a. (a, (GWInfo, [[Expr]])) -> [(a, (Expr, GradInfo))]
go  ((KVar, (GWInfo, [[Expr]])) -> [(KVar, (Expr, GradInfo))])
-> [(KVar, (GWInfo, [[Expr]]))] -> [[(KVar, (Expr, GradInfo))]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(KVar, (GWInfo, [[Expr]]))]
kes))
  where
    go :: (a, (GWInfo, [[Expr]])) -> [(a, (Expr, GradInfo))]
go (a
k, (GWInfo
i, [[Expr]]
es)) = [(a
k, ([Expr] -> Expr
pAnd (GWInfo -> Expr
gexpr GWInfo
iExpr -> [Expr] -> [Expr]
forall a. a -> [a] -> [a]
:[Expr]
e'), GWInfo -> GradInfo
ginfo GWInfo
i)) | [Expr]
e' <- [[Expr]]
es]
    env :: SymEnv
env = Config -> SInfo a -> SymEnv
forall a. Config -> SInfo a -> SymEnv
symbolEnv Config
cfg SInfo a
fi


-------------------------------------------------------------------------------
-- |  Make each gradual appearence unique -------------------------------------
-------------------------------------------------------------------------------
uniquify :: (NFData a, Fixpoint a, Loc a) => SInfo a -> (SInfo a)

uniquify :: SInfo a -> SInfo a
uniquify SInfo a
fi = SInfo a
fi{cm :: HashMap SubcId (SimpC a)
cm = HashMap SubcId (SimpC a)
cm', ws :: HashMap KVar (WfC a)
ws = HashMap KVar (WfC a)
ws', bs :: BindEnv
bs = BindEnv
bs'}
  where
  (HashMap SubcId (SimpC a)
cm', HashMap KVar [(KVar, Maybe SrcSpan)]
km, BindEnv
bs') = BindEnv
-> HashMap SubcId (SimpC a)
-> (HashMap SubcId (SimpC a), HashMap KVar [(KVar, Maybe SrcSpan)],
    BindEnv)
forall a.
(NFData a, Fixpoint a, Loc a) =>
BindEnv
-> HashMap SubcId (SimpC a)
-> (HashMap SubcId (SimpC a), HashMap KVar [(KVar, Maybe SrcSpan)],
    BindEnv)
uniquifyCS (SInfo a -> BindEnv
forall (c :: * -> *) a. GInfo c a -> BindEnv
bs SInfo a
fi) (SInfo a -> HashMap SubcId (SimpC a)
forall (c :: * -> *) a. GInfo c a -> HashMap SubcId (c a)
cm SInfo a
fi)
  ws' :: HashMap KVar (WfC a)
ws'            = HashMap KVar [(KVar, Maybe SrcSpan)]
-> HashMap KVar (WfC a) -> HashMap KVar (WfC a)
forall a.
(NFData a, Fixpoint a) =>
HashMap KVar [(KVar, Maybe SrcSpan)]
-> HashMap KVar (WfC a) -> HashMap KVar (WfC a)
expandWF HashMap KVar [(KVar, Maybe SrcSpan)]
km (SInfo a -> HashMap KVar (WfC a)
forall (c :: * -> *) a. GInfo c a -> HashMap KVar (WfC a)
ws SInfo a
fi)

uniquifyCS :: (NFData a, Fixpoint a, Loc a)
           => BindEnv
           -> M.HashMap SubcId (SimpC a)
           -> (M.HashMap SubcId (SimpC a), M.HashMap KVar [(KVar, Maybe SrcSpan)], BindEnv)
uniquifyCS :: BindEnv
-> HashMap SubcId (SimpC a)
-> (HashMap SubcId (SimpC a), HashMap KVar [(KVar, Maybe SrcSpan)],
    BindEnv)
uniquifyCS BindEnv
bs HashMap SubcId (SimpC a)
cs
  = (HashMap SubcId (SimpC a)
x, HashMap KVar [(KVar, Maybe SrcSpan)]
km, UniqueST -> BindEnv
benv UniqueST
st)
--   = (x, km, mapBindEnv (\i (x,r) -> if i `elem` ubs st then (x, ungrad r) else (x, r)) $ benv st)
  where
    (HashMap SubcId (SimpC a)
x, UniqueST
st) = State UniqueST (HashMap SubcId (SimpC a))
-> UniqueST -> (HashMap SubcId (SimpC a), UniqueST)
forall s a. State s a -> s -> (a, s)
runState (HashMap SubcId (SimpC a)
-> State UniqueST (HashMap SubcId (SimpC a))
forall a. Unique a => a -> UniqueM a
uniq HashMap SubcId (SimpC a)
cs) (BindEnv -> UniqueST
initUniqueST BindEnv
bs)
    km :: HashMap KVar [(KVar, Maybe SrcSpan)]
km      = UniqueST -> HashMap KVar [(KVar, Maybe SrcSpan)]
kmap UniqueST
st
    -- gs      = [x | xs <- M.elems km, (x,_) <- xs]


class Unique a where
   uniq :: a -> UniqueM a

instance Unique a => Unique (M.HashMap SubcId a) where
  uniq :: HashMap SubcId a -> UniqueM (HashMap SubcId a)
uniq HashMap SubcId a
m = [(SubcId, a)] -> HashMap SubcId a
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
M.fromList ([(SubcId, a)] -> HashMap SubcId a)
-> StateT UniqueST Identity [(SubcId, a)]
-> UniqueM (HashMap SubcId a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((SubcId, a) -> StateT UniqueST Identity (SubcId, a))
-> [(SubcId, a)] -> StateT UniqueST Identity [(SubcId, a)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (\(SubcId
i,a
x) -> (SubcId
i,) (a -> (SubcId, a))
-> StateT UniqueST Identity a
-> StateT UniqueST Identity (SubcId, a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> StateT UniqueST Identity a
forall a. Unique a => a -> UniqueM a
uniq a
x) (HashMap SubcId a -> [(SubcId, a)]
forall k v. HashMap k v -> [(k, v)]
M.toList HashMap SubcId a
m)

instance Loc a => Unique (SimpC a) where
  uniq :: SimpC a -> UniqueM (SimpC a)
uniq SimpC a
cs = do
    SrcSpan -> UniqueM ()
updateLoc (SrcSpan -> UniqueM ()) -> SrcSpan -> UniqueM ()
forall a b. (a -> b) -> a -> b
$ a -> SrcSpan
forall a. Loc a => a -> SrcSpan
srcSpan (a -> SrcSpan) -> a -> SrcSpan
forall a b. (a -> b) -> a -> b
$ SimpC a -> a
forall a. SimpC a -> a
_cinfo SimpC a
cs
    Expr
rhs <- Expr -> UniqueM Expr
forall a. Unique a => a -> UniqueM a
uniq (SimpC a -> Expr
forall a. SimpC a -> Expr
_crhs SimpC a
cs)
    IBindEnv
env <- IBindEnv -> UniqueM IBindEnv
forall a. Unique a => a -> UniqueM a
uniq (SimpC a -> IBindEnv
forall a. SimpC a -> IBindEnv
_cenv SimpC a
cs)
    SimpC a -> UniqueM (SimpC a)
forall (m :: * -> *) a. Monad m => a -> m a
return SimpC a
cs{_crhs :: Expr
_crhs = Expr
rhs, _cenv :: IBindEnv
_cenv = IBindEnv
env}

instance Unique IBindEnv where
  uniq :: IBindEnv -> UniqueM IBindEnv
uniq IBindEnv
env = UniqueM IBindEnv -> UniqueM IBindEnv
forall a. UniqueM a -> UniqueM a
withCache ([Int] -> IBindEnv
fromListIBindEnv ([Int] -> IBindEnv)
-> StateT UniqueST Identity [Int] -> UniqueM IBindEnv
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int -> StateT UniqueST Identity Int)
-> [Int] -> StateT UniqueST Identity [Int]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM Int -> StateT UniqueST Identity Int
forall a. Unique a => a -> UniqueM a
uniq (IBindEnv -> [Int]
elemsIBindEnv IBindEnv
env))

instance Unique BindId where
  uniq :: Int -> StateT UniqueST Identity Int
uniq Int
i = do
    BindEnv
bs <- UniqueST -> BindEnv
benv (UniqueST -> BindEnv)
-> StateT UniqueST Identity UniqueST
-> StateT UniqueST Identity BindEnv
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT UniqueST Identity UniqueST
forall s (m :: * -> *). MonadState s m => m s
get
    let (Symbol
x, SortedReft
t) = Int -> BindEnv -> (Symbol, SortedReft)
lookupBindEnv Int
i BindEnv
bs
    UniqueM ()
resetChange
    SortedReft
t' <- SortedReft -> UniqueM SortedReft
forall a. Unique a => a -> UniqueM a
uniq SortedReft
t
    Bool
hasChanged <- UniqueST -> Bool
change (UniqueST -> Bool)
-> StateT UniqueST Identity UniqueST
-> StateT UniqueST Identity Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT UniqueST Identity UniqueST
forall s (m :: * -> *). MonadState s m => m s
get
    if Bool
hasChanged
      then do let (Int
i', BindEnv
bs') = Symbol -> SortedReft -> BindEnv -> (Int, BindEnv)
insertBindEnv Symbol
x SortedReft
t' BindEnv
bs
              Int -> BindEnv -> UniqueM ()
updateBEnv Int
i BindEnv
bs'
              Int -> StateT UniqueST Identity Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
i'
      else Int -> StateT UniqueST Identity Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
i

instance Unique SortedReft where
  uniq :: SortedReft -> UniqueM SortedReft
uniq (RR Sort
s Reft
r) = Sort -> Reft -> SortedReft
RR Sort
s (Reft -> SortedReft)
-> StateT UniqueST Identity Reft -> UniqueM SortedReft
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Reft -> StateT UniqueST Identity Reft
forall a. Unique a => a -> UniqueM a
uniq Reft
r

instance Unique Reft where
  uniq :: Reft -> StateT UniqueST Identity Reft
uniq (Reft (Symbol
x,Expr
e)) = ((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) -> UniqueM Expr -> StateT UniqueST Identity Reft
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr -> UniqueM Expr
forall a. Unique a => a -> UniqueM a
uniq Expr
e

instance Unique Expr where
  uniq :: Expr -> UniqueM Expr
uniq = (Expr -> UniqueM Expr) -> Expr -> UniqueM Expr
forall (m :: * -> *). Monad m => (Expr -> m Expr) -> Expr -> m Expr
mapMExpr Expr -> UniqueM Expr
go
   where
    go :: Expr -> UniqueM Expr
go (PGrad KVar
k Subst
su GradInfo
i Expr
e) = do
      KVar
k'  <- KVar -> UniqueM KVar
freshK KVar
k
      Maybe SrcSpan
src <- UniqueST -> Maybe SrcSpan
uloc (UniqueST -> Maybe SrcSpan)
-> StateT UniqueST Identity UniqueST
-> StateT UniqueST Identity (Maybe SrcSpan)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT UniqueST Identity UniqueST
forall s (m :: * -> *). MonadState s m => m s
get
      Expr -> UniqueM Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> UniqueM Expr) -> Expr -> UniqueM Expr
forall a b. (a -> b) -> a -> b
$ KVar -> Subst -> GradInfo -> Expr -> Expr
PGrad KVar
k' Subst
su (GradInfo
i{gused :: Maybe SrcSpan
gused = Maybe SrcSpan
src}) Expr
e
    go Expr
e              = Expr -> UniqueM Expr
forall (m :: * -> *) a. Monad m => a -> m a
return Expr
e

-------------------------------------------------------------------------------
-- | The Unique Monad ---------------------------------------------------------
-------------------------------------------------------------------------------

type UniqueM = State UniqueST
data UniqueST
  = UniqueST { UniqueST -> SubcId
freshId :: Integer
             , UniqueST -> HashMap KVar [(KVar, Maybe SrcSpan)]
kmap    :: M.HashMap KVar [(KVar, Maybe SrcSpan)]
             , UniqueST -> Bool
change  :: Bool
             , UniqueST -> HashMap KVar KVar
cache   :: M.HashMap KVar KVar
             , UniqueST -> Maybe SrcSpan
uloc    :: Maybe SrcSpan
             , UniqueST -> [Int]
ubs     :: [BindId]
             , UniqueST -> BindEnv
benv    :: BindEnv
             }

updateLoc :: SrcSpan -> UniqueM ()
updateLoc :: SrcSpan -> UniqueM ()
updateLoc SrcSpan
x = (UniqueST -> UniqueST) -> UniqueM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((UniqueST -> UniqueST) -> UniqueM ())
-> (UniqueST -> UniqueST) -> UniqueM ()
forall a b. (a -> b) -> a -> b
$ \UniqueST
s -> UniqueST
s{uloc :: Maybe SrcSpan
uloc = SrcSpan -> Maybe SrcSpan
forall a. a -> Maybe a
Just SrcSpan
x}

withCache :: UniqueM a -> UniqueM a
withCache :: UniqueM a -> UniqueM a
withCache UniqueM a
act = do
  UniqueM ()
emptyCache
  a
a <- UniqueM a
act
  UniqueM ()
emptyCache
  a -> UniqueM a
forall (m :: * -> *) a. Monad m => a -> m a
return a
a

emptyCache :: UniqueM ()
emptyCache :: UniqueM ()
emptyCache = (UniqueST -> UniqueST) -> UniqueM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((UniqueST -> UniqueST) -> UniqueM ())
-> (UniqueST -> UniqueST) -> UniqueM ()
forall a b. (a -> b) -> a -> b
$ \UniqueST
s -> UniqueST
s{cache :: HashMap KVar KVar
cache = HashMap KVar KVar
forall a. Monoid a => a
mempty}

addCache :: KVar -> KVar -> UniqueM ()
addCache :: KVar -> KVar -> UniqueM ()
addCache KVar
k KVar
k' = (UniqueST -> UniqueST) -> UniqueM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((UniqueST -> UniqueST) -> UniqueM ())
-> (UniqueST -> UniqueST) -> UniqueM ()
forall a b. (a -> b) -> a -> b
$ \UniqueST
s -> UniqueST
s{cache :: HashMap KVar KVar
cache = KVar -> KVar -> HashMap KVar KVar -> HashMap KVar KVar
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
M.insert KVar
k KVar
k' (UniqueST -> HashMap KVar KVar
cache UniqueST
s)}

updateBEnv :: BindId -> BindEnv -> UniqueM ()
updateBEnv :: Int -> BindEnv -> UniqueM ()
updateBEnv Int
i BindEnv
bs = (UniqueST -> UniqueST) -> UniqueM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((UniqueST -> UniqueST) -> UniqueM ())
-> (UniqueST -> UniqueST) -> UniqueM ()
forall a b. (a -> b) -> a -> b
$ \UniqueST
s -> UniqueST
s{benv :: BindEnv
benv = BindEnv
bs, ubs :: [Int]
ubs = Int
iInt -> [Int] -> [Int]
forall a. a -> [a] -> [a]
:(UniqueST -> [Int]
ubs UniqueST
s)}

setChange :: UniqueM ()
setChange :: UniqueM ()
setChange = (UniqueST -> UniqueST) -> UniqueM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((UniqueST -> UniqueST) -> UniqueM ())
-> (UniqueST -> UniqueST) -> UniqueM ()
forall a b. (a -> b) -> a -> b
$ \UniqueST
s -> UniqueST
s{change :: Bool
change = Bool
True}

resetChange :: UniqueM ()
resetChange :: UniqueM ()
resetChange = (UniqueST -> UniqueST) -> UniqueM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((UniqueST -> UniqueST) -> UniqueM ())
-> (UniqueST -> UniqueST) -> UniqueM ()
forall a b. (a -> b) -> a -> b
$ \UniqueST
s -> UniqueST
s{change :: Bool
change = Bool
False}

initUniqueST :: BindEnv ->  UniqueST
initUniqueST :: BindEnv -> UniqueST
initUniqueST = SubcId
-> HashMap KVar [(KVar, Maybe SrcSpan)]
-> Bool
-> HashMap KVar KVar
-> Maybe SrcSpan
-> [Int]
-> BindEnv
-> UniqueST
UniqueST SubcId
0 HashMap KVar [(KVar, Maybe SrcSpan)]
forall a. Monoid a => a
mempty Bool
False HashMap KVar KVar
forall a. Monoid a => a
mempty Maybe SrcSpan
forall a. Maybe a
Nothing [Int]
forall a. Monoid a => a
mempty

freshK, freshK' :: KVar -> UniqueM KVar
freshK :: KVar -> UniqueM KVar
freshK KVar
k  = do
  UniqueM ()
setChange
  HashMap KVar KVar
cached <- UniqueST -> HashMap KVar KVar
cache (UniqueST -> HashMap KVar KVar)
-> StateT UniqueST Identity UniqueST
-> StateT UniqueST Identity (HashMap KVar KVar)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT UniqueST Identity UniqueST
forall s (m :: * -> *). MonadState s m => m s
get
  case KVar -> HashMap KVar KVar -> Maybe KVar
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
M.lookup KVar
k HashMap KVar KVar
cached of
    {- OPTIMIZATION: Only create one fresh occurence of ? per constraint environment. -}
    Just KVar
k' -> KVar -> UniqueM KVar
forall (m :: * -> *) a. Monad m => a -> m a
return  KVar
k'
    Maybe KVar
Nothing -> KVar -> UniqueM KVar
freshK' KVar
k

freshK' :: KVar -> UniqueM KVar
freshK' KVar
k = do
  SubcId
i <- UniqueST -> SubcId
freshId (UniqueST -> SubcId)
-> StateT UniqueST Identity UniqueST
-> StateT UniqueST Identity SubcId
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT UniqueST Identity UniqueST
forall s (m :: * -> *). MonadState s m => m s
get
  (UniqueST -> UniqueST) -> UniqueM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((UniqueST -> UniqueST) -> UniqueM ())
-> (UniqueST -> UniqueST) -> UniqueM ()
forall a b. (a -> b) -> a -> b
$ (\UniqueST
s -> UniqueST
s{freshId :: SubcId
freshId = SubcId
i SubcId -> SubcId -> SubcId
forall a. Num a => a -> a -> a
+ SubcId
1})
  let k' :: KVar
k' = Symbol -> KVar
KV (Symbol -> KVar) -> Symbol -> KVar
forall a b. (a -> b) -> a -> b
$ SubcId -> Symbol
gradIntSymbol SubcId
i
  KVar -> KVar -> UniqueM ()
addK KVar
k KVar
k'
  KVar -> KVar -> UniqueM ()
addCache KVar
k KVar
k'
  KVar -> UniqueM KVar
forall (m :: * -> *) a. Monad m => a -> m a
return KVar
k'

addK :: KVar -> KVar -> UniqueM ()
addK :: KVar -> KVar -> UniqueM ()
addK KVar
key KVar
val =
  (UniqueST -> UniqueST) -> UniqueM ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((UniqueST -> UniqueST) -> UniqueM ())
-> (UniqueST -> UniqueST) -> UniqueM ()
forall a b. (a -> b) -> a -> b
$ (\UniqueST
s -> UniqueST
s{kmap :: HashMap KVar [(KVar, Maybe SrcSpan)]
kmap = ([(KVar, Maybe SrcSpan)]
 -> [(KVar, Maybe SrcSpan)] -> [(KVar, Maybe SrcSpan)])
-> KVar
-> [(KVar, Maybe SrcSpan)]
-> HashMap KVar [(KVar, Maybe SrcSpan)]
-> HashMap KVar [(KVar, Maybe SrcSpan)]
forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> k -> v -> HashMap k v -> HashMap k v
M.insertWith [(KVar, Maybe SrcSpan)]
-> [(KVar, Maybe SrcSpan)] -> [(KVar, Maybe SrcSpan)]
forall a. [a] -> [a] -> [a]
(++) KVar
key [(KVar
val, UniqueST -> Maybe SrcSpan
uloc UniqueST
s)] (UniqueST -> HashMap KVar [(KVar, Maybe SrcSpan)]
kmap UniqueST
s)})

-------------------------------------------------------------------------------
-- | expandWF -----------------------------------------------------------------
-------------------------------------------------------------------------------

expandWF :: (NFData a, Fixpoint a)
         => M.HashMap KVar [(KVar, Maybe SrcSpan)]
         -> M.HashMap KVar (WfC a)
         -> M.HashMap KVar (WfC a)
expandWF :: HashMap KVar [(KVar, Maybe SrcSpan)]
-> HashMap KVar (WfC a) -> HashMap KVar (WfC a)
expandWF HashMap KVar [(KVar, Maybe SrcSpan)]
km HashMap KVar (WfC a)
ws
  = [(KVar, WfC a)] -> HashMap KVar (WfC a)
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
M.fromList ([(KVar, WfC a)] -> HashMap KVar (WfC a))
-> [(KVar, WfC a)] -> HashMap KVar (WfC a)
forall a b. (a -> b) -> a -> b
$
       ([(KVar
k, KVar -> Maybe SrcSpan -> WfC a -> WfC a
forall a. KVar -> Maybe SrcSpan -> WfC a -> WfC a
updateKVar KVar
k Maybe SrcSpan
src WfC a
w) | (KVar
i, WfC a
w) <- [(KVar, WfC a)]
gws, (KVar
kw, [(KVar, Maybe SrcSpan)]
ks) <- [(KVar, [(KVar, Maybe SrcSpan)])]
km', KVar
kw KVar -> KVar -> Bool
forall a. Eq a => a -> a -> Bool
== KVar
i, (KVar
k, Maybe SrcSpan
src) <- [(KVar, Maybe SrcSpan)]
ks]
        [(KVar, WfC a)] -> [(KVar, WfC a)] -> [(KVar, WfC a)]
forall a. [a] -> [a] -> [a]
++ [(KVar, WfC a)]
kws)
  where
    ([(KVar, WfC a)]
gws, [(KVar, WfC a)]
kws)       = ((KVar, WfC a) -> Bool)
-> [(KVar, WfC a)] -> ([(KVar, WfC a)], [(KVar, WfC a)])
forall a. (a -> Bool) -> [a] -> ([a], [a])
L.partition (WfC a -> Bool
forall a. WfC a -> Bool
isGWfc (WfC a -> Bool)
-> ((KVar, WfC a) -> WfC a) -> (KVar, WfC a) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (KVar, WfC a) -> WfC a
forall a b. (a, b) -> b
snd) ([(KVar, WfC a)] -> ([(KVar, WfC a)], [(KVar, WfC a)]))
-> [(KVar, WfC a)] -> ([(KVar, WfC a)], [(KVar, WfC a)])
forall a b. (a -> b) -> a -> b
$ HashMap KVar (WfC a) -> [(KVar, WfC a)]
forall k v. HashMap k v -> [(k, v)]
M.toList HashMap KVar (WfC a)
ws
    km' :: [(KVar, [(KVar, Maybe SrcSpan)])]
km'              = HashMap KVar [(KVar, Maybe SrcSpan)]
-> [(KVar, [(KVar, Maybe SrcSpan)])]
forall k v. HashMap k v -> [(k, v)]
M.toList HashMap KVar [(KVar, Maybe SrcSpan)]
km
    updateKVar :: KVar -> Maybe SrcSpan -> WfC a -> WfC a
updateKVar KVar
k Maybe SrcSpan
src WfC a
wfc = WfC a
wfc { wrft :: (Symbol, Sort, KVar)
wrft = (\(Symbol
v,Sort
s,KVar
_) -> (Symbol
v,Sort
s,KVar
k)) ((Symbol, Sort, KVar) -> (Symbol, Sort, KVar))
-> (Symbol, Sort, KVar) -> (Symbol, Sort, KVar)
forall a b. (a -> b) -> a -> b
$ WfC a -> (Symbol, Sort, KVar)
forall a. WfC a -> (Symbol, Sort, KVar)
wrft WfC a
wfc
                               , wloc :: GradInfo
wloc = (WfC a -> GradInfo
forall a. WfC a -> GradInfo
wloc WfC a
wfc){gused :: Maybe SrcSpan
gused = Maybe SrcSpan
src}
                               }

-------------------------------------------------------------------------------
-- |  Substitute Gradual Solution ---------------------------------------------
-------------------------------------------------------------------------------

class Gradual a where
  gsubst :: GSol -> a -> a

instance Gradual Expr where
  gsubst :: GSol -> Expr -> Expr
gsubst (GSol SymEnv
env HashMap KVar (Expr, GradInfo)
m) Expr
e   = ((KVar, Subst) -> Maybe Expr) -> Expr -> Expr
forall t. Visitable t => ((KVar, Subst) -> Maybe Expr) -> t -> t
mapGVars' (\(KVar
k, Subst
_) -> Expr -> Maybe Expr
forall a. a -> Maybe a
Just (Expr -> Maybe Expr -> Expr
forall a. a -> Maybe a -> a
fromMaybe (KVar -> Expr
forall a a. PPrint a => a -> a
err KVar
k) (KVar -> Maybe Expr
mknew KVar
k))) Expr
e
    where
      mknew :: KVar -> Maybe Expr
mknew KVar
k = Located String -> SymEnv -> Maybe Expr -> Maybe Expr
forall a. Elaborate a => Located String -> SymEnv -> a -> a
So.elaborate Located String
"initBGind.mkPred" SymEnv
env (Maybe Expr -> Maybe Expr) -> Maybe Expr -> Maybe Expr
forall a b. (a -> b) -> a -> b
$ (Expr, GradInfo) -> Expr
forall a b. (a, b) -> a
fst ((Expr, GradInfo) -> Expr) -> Maybe (Expr, GradInfo) -> Maybe Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> KVar -> HashMap KVar (Expr, GradInfo) -> Maybe (Expr, GradInfo)
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
M.lookup KVar
k HashMap KVar (Expr, GradInfo)
m
      err :: a -> a
err   a
k = String -> a
forall a. (?callStack::CallStack) => String -> a
errorstar (String
"gradual substitution: Cannot find " String -> ShowS
forall a. [a] -> [a] -> [a]
++ a -> String
forall a. PPrint a => a -> String
showpp a
k)

instance Gradual Reft where
  gsubst :: GSol -> Reft -> Reft
gsubst GSol
su (Reft (Symbol
x, Expr
e)) = (Symbol, Expr) -> Reft
Reft (Symbol
x, GSol -> Expr -> Expr
forall a. Gradual a => GSol -> a -> a
gsubst GSol
su Expr
e)

instance Gradual SortedReft where
  gsubst :: GSol -> SortedReft -> SortedReft
gsubst GSol
su SortedReft
r = SortedReft
r {sr_reft :: Reft
sr_reft = GSol -> Reft -> Reft
forall a. Gradual a => GSol -> a -> a
gsubst GSol
su (SortedReft -> Reft
sr_reft SortedReft
r)}

instance Gradual (SimpC a) where
  gsubst :: GSol -> SimpC a -> SimpC a
gsubst GSol
su SimpC a
c = SimpC a
c {_crhs :: Expr
_crhs = GSol -> Expr -> Expr
forall a. Gradual a => GSol -> a -> a
gsubst GSol
su (SimpC a -> Expr
forall a. SimpC a -> Expr
_crhs SimpC a
c)}

instance Gradual BindEnv where
  gsubst :: GSol -> BindEnv -> BindEnv
gsubst GSol
su = (Int -> (Symbol, SortedReft) -> (Symbol, SortedReft))
-> BindEnv -> BindEnv
mapBindEnv (\Int
_ (Symbol
x, SortedReft
r) -> (Symbol
x, GSol -> SortedReft -> SortedReft
forall a. Gradual a => GSol -> a -> a
gsubst GSol
su SortedReft
r))

instance Gradual v => Gradual (M.HashMap k v) where
  gsubst :: GSol -> HashMap k v -> HashMap k v
gsubst GSol
su = (v -> v) -> HashMap k v -> HashMap k v
forall v1 v2 k. (v1 -> v2) -> HashMap k v1 -> HashMap k v2
M.map (GSol -> v -> v
forall a. Gradual a => GSol -> a -> a
gsubst GSol
su)

instance Gradual (SInfo a) where
  gsubst :: GSol -> SInfo a -> SInfo a
gsubst GSol
su SInfo a
fi = SInfo a
fi { bs :: BindEnv
bs = GSol -> BindEnv -> BindEnv
forall a. Gradual a => GSol -> a -> a
gsubst GSol
su (SInfo a -> BindEnv
forall (c :: * -> *) a. GInfo c a -> BindEnv
bs SInfo a
fi)
                    , cm :: HashMap SubcId (SimpC a)
cm = GSol -> HashMap SubcId (SimpC a) -> HashMap SubcId (SimpC a)
forall a. Gradual a => GSol -> a -> a
gsubst GSol
su (SInfo a -> HashMap SubcId (SimpC a)
forall (c :: * -> *) a. GInfo c a -> HashMap SubcId (c a)
cm SInfo a
fi)
                    }