{-# LANGUAGE CPP                        #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE NoMonomorphismRestriction  #-}
{-# LANGUAGE OverloadedStrings          #-}
{-# LANGUAGE UndecidableInstances       #-}
{-# LANGUAGE MultiParamTypeClasses      #-}
{-# LANGUAGE GADTs                      #-}
{-# LANGUAGE TupleSections              #-}

{-# OPTIONS_GHC -Wno-name-shadowing     #-}

-- | 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)
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
forall v. v -> ExprV v
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 :: forall a.
(NFData a, Fixpoint a, Show a) =>
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
forall v. (Fixpoint v, Ord v) => ListNE (ExprV v) -> ExprV v
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 :: forall a. (NFData a, Fixpoint a, Loc a) => SInfo a -> SInfo a
uniquify SInfo a
fi = SInfo a
fi{cm = cm', ws = ws', bs = bs'}
  where
  (HashMap SubcId (SimpC a)
cm', HashMap KVar [(KVar, Maybe SrcSpan)]
km, BindEnv a
bs') = BindEnv a
-> HashMap SubcId (SimpC a)
-> (HashMap SubcId (SimpC a), HashMap KVar [(KVar, Maybe SrcSpan)],
    BindEnv a)
forall a.
(NFData a, Fixpoint a, Loc a) =>
BindEnv a
-> HashMap SubcId (SimpC a)
-> (HashMap SubcId (SimpC a), HashMap KVar [(KVar, Maybe SrcSpan)],
    BindEnv a)
uniquifyCS (SInfo a -> BindEnv a
forall (c :: * -> *) a. GInfo c a -> BindEnv a
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 a
           -> M.HashMap SubcId (SimpC a)
           -> (M.HashMap SubcId (SimpC a), M.HashMap KVar [(KVar, Maybe SrcSpan)], BindEnv a)
uniquifyCS :: forall a.
(NFData a, Fixpoint a, Loc a) =>
BindEnv a
-> HashMap SubcId (SimpC a)
-> (HashMap SubcId (SimpC a), HashMap KVar [(KVar, Maybe SrcSpan)],
    BindEnv a)
uniquifyCS BindEnv a
bs HashMap SubcId (SimpC a)
cs
  = (HashMap SubcId (SimpC a)
x, HashMap KVar [(KVar, Maybe SrcSpan)]
km, UniqueST a -> BindEnv a
forall a. UniqueST a -> BindEnv a
benv UniqueST a
st)
  where
    (HashMap SubcId (SimpC a)
x, UniqueST a
st) = State (UniqueST a) (HashMap SubcId (SimpC a))
-> UniqueST a -> (HashMap SubcId (SimpC a), UniqueST a)
forall s a. State s a -> s -> (a, s)
runState (HashMap SubcId (SimpC a)
-> State (UniqueST a) (HashMap SubcId (SimpC a))
forall ann a. Unique ann a => a -> UniqueM ann a
uniq HashMap SubcId (SimpC a)
cs) (BindEnv a -> UniqueST a
forall a. BindEnv a -> UniqueST a
initUniqueST BindEnv a
bs)
    km :: HashMap KVar [(KVar, Maybe SrcSpan)]
km      = UniqueST a -> HashMap KVar [(KVar, Maybe SrcSpan)]
forall a. UniqueST a -> HashMap KVar [(KVar, Maybe SrcSpan)]
kmap UniqueST a
st


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

instance Unique ann a => Unique ann (M.HashMap SubcId a) where
  uniq :: HashMap SubcId a -> UniqueM ann (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 ann) Identity [(SubcId, a)]
-> UniqueM ann (HashMap SubcId a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((SubcId, a) -> StateT (UniqueST ann) Identity (SubcId, a))
-> [(SubcId, a)] -> StateT (UniqueST ann) Identity [(SubcId, a)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (\(SubcId
i,a
x) -> (SubcId
i,) (a -> (SubcId, a))
-> StateT (UniqueST ann) Identity a
-> StateT (UniqueST ann) Identity (SubcId, a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> StateT (UniqueST ann) Identity a
forall ann a. Unique ann a => a -> UniqueM ann 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 a (SimpC a) where
  uniq :: SimpC a -> UniqueM a (SimpC a)
uniq SimpC a
cs = do
    SrcSpan -> UniqueM a ()
forall ann. SrcSpan -> UniqueM ann ()
updateLoc (SrcSpan -> UniqueM a ()) -> SrcSpan -> UniqueM a ()
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 a Expr
forall ann a. Unique ann a => a -> UniqueM ann a
uniq (SimpC a -> Expr
forall a. SimpC a -> Expr
_crhs SimpC a
cs)
    IBindEnv
env <- IBindEnv -> UniqueM a IBindEnv
forall ann a. Unique ann a => a -> UniqueM ann a
uniq (SimpC a -> IBindEnv
forall a. SimpC a -> IBindEnv
_cenv SimpC a
cs)
    SimpC a -> UniqueM a (SimpC a)
forall a. a -> StateT (UniqueST a) Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return SimpC a
cs{_crhs = rhs, _cenv = env}

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

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

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

instance Unique ann Reft where
  uniq :: Reft -> UniqueM ann Reft
uniq (Reft (Symbol
x,Expr
e)) = (Symbol, Expr) -> Reft
forall v. (Symbol, ExprV v) -> ReftV v
Reft ((Symbol, Expr) -> Reft)
-> (Expr -> (Symbol, Expr)) -> Expr -> Reft
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Symbol
x,) (Expr -> Reft)
-> StateT (UniqueST ann) Identity Expr -> UniqueM ann Reft
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr -> StateT (UniqueST ann) Identity Expr
forall ann a. Unique ann a => a -> UniqueM ann a
uniq Expr
e

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

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

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

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

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

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

addCache :: KVar -> KVar -> UniqueM ann ()
addCache :: forall ann. KVar -> KVar -> UniqueM ann ()
addCache KVar
k KVar
k' = (UniqueST ann -> UniqueST ann) -> StateT (UniqueST ann) Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((UniqueST ann -> UniqueST ann)
 -> StateT (UniqueST ann) Identity ())
-> (UniqueST ann -> UniqueST ann)
-> StateT (UniqueST ann) Identity ()
forall a b. (a -> b) -> a -> b
$ \UniqueST ann
s -> UniqueST ann
s{cache = M.insert k k' (cache s)}

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

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

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

initUniqueST :: BindEnv a ->  UniqueST a
initUniqueST :: forall a. BindEnv a -> UniqueST a
initUniqueST = SubcId
-> HashMap KVar [(KVar, Maybe SrcSpan)]
-> Bool
-> HashMap KVar KVar
-> Maybe SrcSpan
-> [Int]
-> BindEnv a
-> UniqueST a
forall a.
SubcId
-> HashMap KVar [(KVar, Maybe SrcSpan)]
-> Bool
-> HashMap KVar KVar
-> Maybe SrcSpan
-> [Int]
-> BindEnv a
-> UniqueST a
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 ann KVar
freshK :: forall ann. KVar -> UniqueM ann KVar
freshK KVar
k  = do
  UniqueM ann ()
forall ann. UniqueM ann ()
setChange
  HashMap KVar KVar
cached <- UniqueST ann -> HashMap KVar KVar
forall a. UniqueST a -> HashMap KVar KVar
cache (UniqueST ann -> HashMap KVar KVar)
-> StateT (UniqueST ann) Identity (UniqueST ann)
-> StateT (UniqueST ann) Identity (HashMap KVar KVar)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT (UniqueST ann) Identity (UniqueST ann)
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 ann KVar
forall a. a -> StateT (UniqueST ann) Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return  KVar
k'
    Maybe KVar
Nothing -> KVar -> UniqueM ann KVar
forall ann. KVar -> UniqueM ann KVar
freshK' KVar
k

freshK' :: forall ann. KVar -> UniqueM ann KVar
freshK' KVar
k = do
  SubcId
i <- UniqueST ann -> SubcId
forall a. UniqueST a -> SubcId
freshId (UniqueST ann -> SubcId)
-> StateT (UniqueST ann) Identity (UniqueST ann)
-> StateT (UniqueST ann) Identity SubcId
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StateT (UniqueST ann) Identity (UniqueST ann)
forall s (m :: * -> *). MonadState s m => m s
get
  (UniqueST ann -> UniqueST ann) -> StateT (UniqueST ann) Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (\UniqueST ann
s -> UniqueST ann
s{freshId = i + 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 -> StateT (UniqueST ann) Identity ()
forall ann. KVar -> KVar -> UniqueM ann ()
addK KVar
k KVar
k'
  KVar -> KVar -> StateT (UniqueST ann) Identity ()
forall ann. KVar -> KVar -> UniqueM ann ()
addCache KVar
k KVar
k'
  KVar -> UniqueM ann KVar
forall a. a -> StateT (UniqueST ann) Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return KVar
k'

addK :: KVar -> KVar -> UniqueM ann ()
addK :: forall ann. KVar -> KVar -> UniqueM ann ()
addK KVar
key KVar
val =
  (UniqueST ann -> UniqueST ann) -> StateT (UniqueST ann) Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (\UniqueST ann
s -> UniqueST ann
s{kmap = M.insertWith (++) key [(val, uloc s)] (kmap s)})

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

expandWF :: (NFData a, Fixpoint a)
         => M.HashMap KVar [(KVar, Maybe SrcSpan)]
         -> M.HashMap KVar (WfC a)
         -> M.HashMap KVar (WfC a)
expandWF :: 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 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
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 = let 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 in
      case WfC a
wfc of
        GWfC{} -> WfC a
wfc { wrft = wrft', wloc = (wloc wfc){gused = src} }
        WfC{}  -> WfC a
wfc { wrft = wrft' }

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

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

instance Gradual Expr where
  gsubst :: ElabFlags -> GSol -> Expr -> Expr
gsubst ElabFlags
ef (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 = ElabParam -> Maybe Expr -> Maybe Expr
forall a. Elaborate a => ElabParam -> a -> a
So.elaborate (ElabFlags -> Located String -> SymEnv -> ElabParam
So.ElabParam ElabFlags
ef 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 :: ElabFlags -> GSol -> Reft -> Reft
gsubst ElabFlags
ef GSol
su (Reft (Symbol
x, Expr
e)) = (Symbol, Expr) -> Reft
forall v. (Symbol, ExprV v) -> ReftV v
Reft (Symbol
x, ElabFlags -> GSol -> Expr -> Expr
forall a. Gradual a => ElabFlags -> GSol -> a -> a
gsubst ElabFlags
ef GSol
su Expr
e)

instance Gradual SortedReft where
  gsubst :: ElabFlags -> GSol -> SortedReft -> SortedReft
gsubst ElabFlags
ef GSol
su SortedReft
r = SortedReft
r {sr_reft = gsubst ef su (sr_reft r)}

instance Gradual (SimpC a) where
  gsubst :: ElabFlags -> GSol -> SimpC a -> SimpC a
gsubst ElabFlags
ef GSol
su SimpC a
c = SimpC a
c {_crhs = gsubst ef su (_crhs c)}

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

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

instance Gradual (SInfo a) where
  gsubst :: ElabFlags -> GSol -> SInfo a -> SInfo a
gsubst ElabFlags
ef GSol
su SInfo a
fi = SInfo a
fi { bs = gsubst ef su (bs fi)
                        , cm = gsubst ef su (cm fi)
                        }