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

module Gradual.Uniquify (uniquify) where

import Language.Fixpoint.Types.Refinements
import Language.Fixpoint.Types.Constraints
import Language.Fixpoint.Types.PrettyPrint
import Language.Fixpoint.Types.Environments
import Language.Fixpoint.Types.Visitor
import Language.Fixpoint.Types.Spans
import Language.Fixpoint.Types.Names        (gradIntSymbol)

import Control.DeepSeq

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

import Control.Monad.State.Lazy

import Gradual.Types (GSpan)

-------------------------------------------------------------------------------
-- |  Make each gradual appearance unique -------------------------------------
-------------------------------------------------------------------------------

uniquify :: (NFData a, Fixpoint a, Loc a) => SInfo a -> (GSpan, SInfo a)

uniquify :: SInfo a -> (GSpan, SInfo a)
uniquify SInfo a
fi = (GSpan
km', 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', GSpan
km, BindEnv
bs') = BindEnv
-> HashMap SubcId (SimpC a)
-> (HashMap SubcId (SimpC a), GSpan, BindEnv)
forall a.
(NFData a, Fixpoint a, Loc a) =>
BindEnv
-> HashMap SubcId (SimpC a)
-> (HashMap SubcId (SimpC a), GSpan, 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'            = GSpan -> HashMap KVar (WfC a) -> HashMap KVar (WfC a)
forall a.
(NFData a, Fixpoint a) =>
GSpan -> HashMap KVar (WfC a) -> HashMap KVar (WfC a)
expandWF GSpan
km (SInfo a -> HashMap KVar (WfC a)
forall (c :: * -> *) a. GInfo c a -> HashMap KVar (WfC a)
ws SInfo a
fi)
  km' :: GSpan
km'            = GSpan -> GSpan -> GSpan
forall k v.
(Eq k, Hashable k) =>
HashMap k v -> HashMap k v -> HashMap k v
M.union GSpan
km (GSpan -> GSpan) -> GSpan -> GSpan
forall a b. (a -> b) -> a -> b
$ (WfC a -> [(KVar, Maybe SrcSpan)]) -> HashMap KVar (WfC a) -> GSpan
forall v1 v2 k. (v1 -> v2) -> HashMap k v1 -> HashMap k v2
M.map ([(KVar, Maybe SrcSpan)] -> WfC a -> [(KVar, Maybe SrcSpan)]
forall a b. a -> b -> a
const []) (HashMap KVar (WfC a) -> GSpan) -> HashMap KVar (WfC a) -> GSpan
forall a b. (a -> b) -> a -> b
$ (WfC a -> Bool) -> HashMap KVar (WfC a) -> HashMap KVar (WfC a)
forall v k. (v -> Bool) -> HashMap k v -> HashMap k v
M.filter WfC a -> Bool
forall a. HasGradual a => a -> Bool
isGradual (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), GSpan, BindEnv)
uniquifyCS :: BindEnv
-> HashMap SubcId (SimpC a)
-> (HashMap SubcId (SimpC a), GSpan, BindEnv)
uniquifyCS BindEnv
bs HashMap SubcId (SimpC a)
cs = (HashMap SubcId (SimpC a)
x, GSpan
km, UniqueST -> BindEnv
benv UniqueST
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 :: GSpan
km      = UniqueST -> GSpan
kmap UniqueST
st

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 ([BindId] -> IBindEnv
fromListIBindEnv ([BindId] -> IBindEnv)
-> StateT UniqueST Identity [BindId] -> UniqueM IBindEnv
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (BindId -> StateT UniqueST Identity BindId)
-> [BindId] -> StateT UniqueST Identity [BindId]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM BindId -> StateT UniqueST Identity BindId
forall a. Unique a => a -> UniqueM a
uniq (IBindEnv -> [BindId]
elemsIBindEnv IBindEnv
env))

instance Unique BindId where
  uniq :: BindId -> StateT UniqueST Identity BindId
uniq BindId
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) = BindId -> BindEnv -> (Symbol, SortedReft)
lookupBindEnv BindId
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 (BindId
i', BindEnv
bs') = Symbol -> SortedReft -> BindEnv -> (BindId, BindEnv)
insertBindEnv Symbol
x SortedReft
t' BindEnv
bs
              BindId -> BindEnv -> UniqueM ()
updateBEnv BindId
i BindEnv
bs'
              BindId -> StateT UniqueST Identity BindId
forall (m :: * -> *) a. Monad m => a -> m a
return BindId
i'
      else BindId -> StateT UniqueST Identity BindId
forall (m :: * -> *) a. Monad m => a -> m a
return BindId
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 -> GSpan
kmap    :: GSpan
             , UniqueST -> Bool
change  :: Bool
             , UniqueST -> HashMap KVar KVar
cache   :: M.HashMap KVar KVar
             , UniqueST -> Maybe SrcSpan
uloc    :: Maybe SrcSpan
             , UniqueST -> [BindId]
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}


updateBEnv :: BindId -> BindEnv -> UniqueM ()
updateBEnv :: BindId -> BindEnv -> UniqueM ()
updateBEnv BindId
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 :: [BindId]
ubs = BindId
iBindId -> [BindId] -> [BindId]
forall a. a -> [a] -> [a]
:(UniqueST -> [BindId]
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
-> GSpan
-> Bool
-> HashMap KVar KVar
-> Maybe SrcSpan
-> [BindId]
-> BindEnv
-> UniqueST
UniqueST SubcId
0 GSpan
forall a. Monoid a => a
mempty Bool
False HashMap KVar KVar
forall a. Monoid a => a
mempty Maybe SrcSpan
forall a. Maybe a
Nothing [BindId]
forall a. Monoid a => a
mempty

freshK, 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 occurrence 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

-- M.HashMap KVar [(KVar, Maybe SrcSpan)]

freshK'' :: KVar -> UniqueM KVar
freshK'' KVar
k = do 
  GSpan
curr <- UniqueST -> GSpan
kmap (UniqueST -> GSpan)
-> StateT UniqueST Identity UniqueST
-> StateT UniqueST Identity GSpan
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 
  Maybe SrcSpan
loc  <- 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 
  case KVar -> GSpan -> Maybe [(KVar, Maybe SrcSpan)]
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
M.lookup KVar
k GSpan
curr of 
    Maybe [(KVar, Maybe SrcSpan)]
Nothing -> KVar -> UniqueM KVar
freshK' KVar
k 
    Just [(KVar, Maybe SrcSpan)]
xs -> do Maybe KVar
spans <- KVar
-> Maybe SrcSpan -> [(KVar, Maybe SrcSpan)] -> UniqueM (Maybe KVar)
existingSpan KVar
k Maybe SrcSpan
loc [(KVar, Maybe SrcSpan)]
xs 
                  case Maybe KVar
spans of 
                    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 


{-
Interesting Relative Positions

           [.......]
Ouside     |  [..]
Inside   [............]
Smaller    [..............] 
Greater      [....]

-}

data RelativePos a
  = Outside  a 
  | Inside   a 
  | Smaller  a 
  | Greater  a
  | NoRelative

existingSpan :: KVar -> Maybe SrcSpan -> [(KVar, Maybe SrcSpan)] -> UniqueM (Maybe KVar)
existingSpan :: KVar
-> Maybe SrcSpan -> [(KVar, Maybe SrcSpan)] -> UniqueM (Maybe KVar)
existingSpan KVar
_ Maybe SrcSpan
Nothing [(KVar, Maybe SrcSpan)]
_ = Maybe KVar -> UniqueM (Maybe KVar)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe KVar
forall a. Maybe a
Nothing 
existingSpan KVar
k (Just SrcSpan
span) [(KVar, Maybe SrcSpan)]
kspans
  = case [(KVar, Maybe SrcSpan)] -> RelativePos (KVar, SrcSpan)
go [(KVar, Maybe SrcSpan)]
kspans of 
      Inside  (KVar
k',SrcSpan
_)  -> KVar -> KVar -> SrcSpan -> UniqueM ()
updateK KVar
k KVar
k' SrcSpan
span                 UniqueM () -> UniqueM (Maybe KVar) -> UniqueM (Maybe KVar)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (Maybe KVar -> UniqueM (Maybe KVar)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe KVar -> UniqueM (Maybe KVar))
-> Maybe KVar -> UniqueM (Maybe KVar)
forall a b. (a -> b) -> a -> b
$ KVar -> Maybe KVar
forall a. a -> Maybe a
Just KVar
k')
      Outside (KVar
k',SrcSpan
sp) -> SrcSpan -> UniqueM ()
updateLoc SrcSpan
sp                      UniqueM () -> UniqueM (Maybe KVar) -> UniqueM (Maybe KVar)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (Maybe KVar -> UniqueM (Maybe KVar)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe KVar -> UniqueM (Maybe KVar))
-> Maybe KVar -> UniqueM (Maybe KVar)
forall a b. (a -> b) -> a -> b
$ KVar -> Maybe KVar
forall a. a -> Maybe a
Just KVar
k')
      Smaller (KVar
k',SrcSpan
sp) -> KVar -> KVar -> SrcSpan -> UniqueM ()
updateK KVar
k KVar
k' (SrcSpan
sp SrcSpan -> SrcSpan -> SrcSpan
`spanDiff` SrcSpan
span) UniqueM () -> UniqueM (Maybe KVar) -> UniqueM (Maybe KVar)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (Maybe KVar -> UniqueM (Maybe KVar)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe KVar -> UniqueM (Maybe KVar))
-> Maybe KVar -> UniqueM (Maybe KVar)
forall a b. (a -> b) -> a -> b
$ KVar -> Maybe KVar
forall a. a -> Maybe a
Just KVar
k')
      Greater (KVar
k',SrcSpan
sp) -> SrcSpan -> UniqueM ()
updateLoc    (SrcSpan
sp SrcSpan -> SrcSpan -> SrcSpan
`spanDiff` SrcSpan
span) UniqueM () -> UniqueM (Maybe KVar) -> UniqueM (Maybe KVar)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (Maybe KVar -> UniqueM (Maybe KVar)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe KVar -> UniqueM (Maybe KVar))
-> Maybe KVar -> UniqueM (Maybe KVar)
forall a b. (a -> b) -> a -> b
$ KVar -> Maybe KVar
forall a. a -> Maybe a
Just KVar
k')
      RelativePos (KVar, SrcSpan)
NoRelative    -> Maybe KVar -> UniqueM (Maybe KVar)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe KVar
forall a. Maybe a
Nothing

  where
    go :: [(KVar, Maybe SrcSpan)] -> RelativePos (KVar, SrcSpan)
go []                    = RelativePos (KVar, SrcSpan)
forall a. RelativePos a
NoRelative
    go ((KVar
_,Maybe SrcSpan
Nothing):[(KVar, Maybe SrcSpan)]
sps)     = [(KVar, Maybe SrcSpan)] -> RelativePos (KVar, SrcSpan)
go [(KVar, Maybe SrcSpan)]
sps
    go ((KVar
x,Just SrcSpan
sp):[(KVar, Maybe SrcSpan)]
sps) 
      | SrcSpan
span SrcSpan -> SrcSpan -> Bool
`inSpan` SrcSpan
sp    = (KVar, SrcSpan) -> RelativePos (KVar, SrcSpan)
forall a. a -> RelativePos a
Inside  (KVar
x, SrcSpan
sp) 
      | SrcSpan
sp SrcSpan -> SrcSpan -> Bool
`inSpan` SrcSpan
span    = (KVar, SrcSpan) -> RelativePos (KVar, SrcSpan)
forall a. a -> RelativePos a
Outside (KVar
x, SrcSpan
sp) 
      | SrcSpan
span SrcSpan -> SrcSpan -> Bool
`isSmaller` SrcSpan
sp = (KVar, SrcSpan) -> RelativePos (KVar, SrcSpan)
forall a. a -> RelativePos a
Smaller (KVar
x, SrcSpan
sp)
      | SrcSpan
sp SrcSpan -> SrcSpan -> Bool
`isSmaller` SrcSpan
span = (KVar, SrcSpan) -> RelativePos (KVar, SrcSpan)
forall a. a -> RelativePos a
Greater (KVar
x, SrcSpan
sp)
      | Bool
otherwise           = [(KVar, Maybe SrcSpan)] -> RelativePos (KVar, SrcSpan)
go [(KVar, Maybe SrcSpan)]
sps 



spanDiff :: SrcSpan -> SrcSpan -> SrcSpan
spanDiff :: SrcSpan -> SrcSpan -> SrcSpan
spanDiff (SS SourcePos
_ SourcePos
end1) (SS SourcePos
_ SourcePos
end2) = SourcePos -> SourcePos -> SrcSpan
SS ((SourceName, BindId, BindId) -> SourcePos
toSourcePos (SourceName
f,BindId
l,BindId
cBindId -> BindId -> BindId
forall a. Num a => a -> a -> a
-BindId
1)) SourcePos
end2
  where (SourceName
f,BindId
l,BindId
c) = SourcePos -> (SourceName, BindId, BindId)
sourcePosElts SourcePos
end1

inSpan :: SrcSpan -> SrcSpan -> Bool 
inSpan :: SrcSpan -> SrcSpan -> Bool
inSpan (SS SourcePos
start1 SourcePos
end1) (SS SourcePos
start2 SourcePos
end2) 
  = (SourceName, BindId, BindId)
start1' (SourceName, BindId, BindId)
-> (SourceName, BindId, BindId) -> Bool
forall a. Ord a => a -> a -> Bool
> (SourceName, BindId, BindId)
start2' Bool -> Bool -> Bool
&& (SourceName, BindId, BindId)
end1' (SourceName, BindId, BindId)
-> (SourceName, BindId, BindId) -> Bool
forall a. Ord a => a -> a -> Bool
<= (SourceName, BindId, BindId)
end2'
  where
    [(SourceName, BindId, BindId)
start1',(SourceName, BindId, BindId)
start2',(SourceName, BindId, BindId)
end1',(SourceName, BindId, BindId)
end2'] = SourcePos -> (SourceName, BindId, BindId)
sourcePosElts (SourcePos -> (SourceName, BindId, BindId))
-> [SourcePos] -> [(SourceName, BindId, BindId)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [SourcePos
start1,SourcePos
start2,SourcePos
end1,SourcePos
end2]  


isSmaller :: SrcSpan -> SrcSpan -> Bool 
isSmaller :: SrcSpan -> SrcSpan -> Bool
isSmaller (SS SourcePos
start1 SourcePos
end1) (SS SourcePos
start2 SourcePos
end2) 
  = (SourceName, BindId, BindId)
start1' (SourceName, BindId, BindId)
-> (SourceName, BindId, BindId) -> Bool
forall a. Eq a => a -> a -> Bool
== (SourceName, BindId, BindId)
start2' Bool -> Bool -> Bool
&& (SourceName, BindId, BindId)
end1' (SourceName, BindId, BindId)
-> (SourceName, BindId, BindId) -> Bool
forall a. Ord a => a -> a -> Bool
<= (SourceName, BindId, BindId)
end2'
  where
    [(SourceName, BindId, BindId)
start1',(SourceName, BindId, BindId)
start2',(SourceName, BindId, BindId)
end1',(SourceName, BindId, BindId)
end2'] = SourcePos -> (SourceName, BindId, BindId)
sourcePosElts (SourcePos -> (SourceName, BindId, BindId))
-> [SourcePos] -> [(SourceName, BindId, BindId)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [SourcePos
start1,SourcePos
start2,SourcePos
end1,SourcePos
end2]  


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'


updateK :: KVar -> KVar -> SrcSpan -> UniqueM ()
updateK :: KVar -> KVar -> SrcSpan -> UniqueM ()
updateK KVar
k KVar
k' SrcSpan
sp = 
  (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 :: GSpan
kmap = (KVar -> [(KVar, Maybe SrcSpan)] -> [(KVar, Maybe SrcSpan)])
-> GSpan -> GSpan
forall k v1 v2. (k -> v1 -> v2) -> HashMap k v1 -> HashMap k v2
M.mapWithKey KVar -> [(KVar, Maybe SrcSpan)] -> [(KVar, Maybe SrcSpan)]
f (UniqueST -> GSpan
kmap UniqueST
s)})
  where
    f :: KVar -> [(KVar, Maybe SrcSpan)] -> [(KVar, Maybe SrcSpan)]
f KVar
key [(KVar, Maybe SrcSpan)]
val 
      | KVar
key KVar -> KVar -> Bool
forall a. Eq a => a -> a -> Bool
== KVar
k  = (\(KVar
kk,Maybe SrcSpan
vv) -> if KVar
kk KVar -> KVar -> Bool
forall a. Eq a => a -> a -> Bool
== KVar
k' then (KVar
kk,SrcSpan -> Maybe SrcSpan
forall a. a -> Maybe a
Just SrcSpan
sp) else (KVar
kk,Maybe SrcSpan
vv)) ((KVar, Maybe SrcSpan) -> (KVar, Maybe SrcSpan))
-> [(KVar, Maybe SrcSpan)] -> [(KVar, Maybe SrcSpan)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(KVar, Maybe SrcSpan)]
val 
      | Bool
otherwise = [(KVar, Maybe SrcSpan)]
val 

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 :: GSpan
kmap = ([(KVar, Maybe SrcSpan)]
 -> [(KVar, Maybe SrcSpan)] -> [(KVar, Maybe SrcSpan)])
-> KVar -> [(KVar, Maybe SrcSpan)] -> GSpan -> GSpan
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 -> GSpan
kmap UniqueST
s)})

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

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

expandWF :: (NFData a, Fixpoint a)
         => GSpan
         -> M.HashMap KVar (WfC a)
         -> M.HashMap KVar (WfC a)
expandWF :: GSpan -> HashMap KVar (WfC a) -> HashMap KVar (WfC a)
expandWF GSpan
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'              = GSpan -> [(KVar, [(KVar, Maybe SrcSpan)])]
forall k v. HashMap k v -> [(k, v)]
M.toList GSpan
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}
                               }