{-# LANGUAGE ScopedTypeVariables       #-}
{-# LANGUAGE NoMonomorphismRestriction #-}
{-# LANGUAGE TypeSynonymInstances      #-}
{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE FlexibleInstances         #-}
{-# LANGUAGE TupleSections             #-}
{-# LANGUAGE BangPatterns              #-}
{-# LANGUAGE PatternGuards             #-}
{-# LANGUAGE DeriveFunctor             #-}
{-# LANGUAGE MultiParamTypeClasses     #-}
{-# LANGUAGE OverloadedStrings         #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ImplicitParams            #-}
{-# LANGUAGE PartialTypeSignatures     #-}

-- | This module defines the representation for Environments needed
--   during constraint generation.

module Language.Haskell.Liquid.Constraint.Env (

  -- * Insert
    (+++=)
  -- , (++=)
  , (+=)
  , extendEnvWithVV
  , addBinders
  , addSEnv
  , addEEnv
  , (-=)
  , globalize

  -- * Construction
  , fromListREnv
  , toListREnv
  , insertREnv -- TODO: remove this ASAP

  -- * Query
  , localBindsOfType
  , lookupREnv
  , (?=)

  -- * Pruning refinements (TODO: move!)
 , rTypeSortedReft'

  -- * Extend CGEnv
 , setLocation, setBind, setRecs, setTRec

  -- * Lookup CGEnv
 , getLocation

) where


-- import Name (getSrcSpan)
import Prelude hiding (error)
import CoreSyn (CoreExpr)
import SrcLoc
import Var
-- import Outputable
-- import FastString (fsLit)
import Control.Monad.State

-- import           GHC.Err.Located hiding (error)
import           GHC.Stack

import           Control.Arrow           (first)
import           Data.Maybe              -- (fromMaybe)
import qualified Data.List               as L
import qualified Data.HashSet            as S
import qualified Data.HashMap.Strict     as M
import qualified Language.Fixpoint.Types as F


import           Language.Fixpoint.SortCheck (pruneUnsortedReft)



import           Language.Haskell.Liquid.Types.RefType
import qualified Language.Haskell.Liquid.GHC.SpanStack as Sp
import           Language.Haskell.Liquid.Types            hiding (binds, Loc, loc, freeTyVars, Def)
import           Language.Haskell.Liquid.Constraint.Types
import           Language.Haskell.Liquid.Constraint.Fresh ()
import           Language.Haskell.Liquid.Transforms.RefSplit
import qualified Language.Haskell.Liquid.UX.CTags       as Tg

-- import Debug.Trace (trace)
--------------------------------------------------------------------------------
-- | Refinement Type Environments ----------------------------------------------
--------------------------------------------------------------------------------

-- updREnvLocal :: REnv -> (_ -> _) -> REnv
updREnvLocal :: REnv
             -> (M.HashMap F.Symbol SpecType -> M.HashMap F.Symbol SpecType)
             -> REnv
updREnvLocal :: REnv
-> (HashMap Symbol SpecType -> HashMap Symbol SpecType) -> REnv
updREnvLocal REnv
rE HashMap Symbol SpecType -> HashMap Symbol SpecType
f      = REnv
rE { reLocal :: HashMap Symbol SpecType
reLocal = HashMap Symbol SpecType -> HashMap Symbol SpecType
f (REnv -> HashMap Symbol SpecType
forall t. AREnv t -> HashMap Symbol t
reLocal REnv
rE) }

-- RJ: REnv-Split-Bug?
filterREnv :: (SpecType -> Bool) -> REnv -> REnv
filterREnv :: (SpecType -> Bool) -> REnv -> REnv
filterREnv SpecType -> Bool
f REnv
rE        = REnv
rE REnv
-> (HashMap Symbol SpecType -> HashMap Symbol SpecType) -> REnv
`updREnvLocal` ((SpecType -> Bool)
-> HashMap Symbol SpecType -> HashMap Symbol SpecType
forall v k. (v -> Bool) -> HashMap k v -> HashMap k v
M.filter SpecType -> Bool
f)

fromListREnv :: [(F.Symbol, SpecType)] -> [(F.Symbol, SpecType)] -> REnv
fromListREnv :: [(Symbol, SpecType)] -> [(Symbol, SpecType)] -> REnv
fromListREnv [(Symbol, SpecType)]
gXts [(Symbol, SpecType)]
lXts = REnv :: forall t. HashMap Symbol t -> HashMap Symbol t -> AREnv t
REnv
  { reGlobal :: HashMap Symbol SpecType
reGlobal = [(Symbol, SpecType)] -> HashMap Symbol SpecType
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
M.fromList [(Symbol, SpecType)]
gXts
  , reLocal :: HashMap Symbol SpecType
reLocal  = [(Symbol, SpecType)] -> HashMap Symbol SpecType
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
M.fromList [(Symbol, SpecType)]
lXts
  }

-- RJ: REnv-Split-Bug?
deleteREnv :: F.Symbol -> REnv -> REnv
deleteREnv :: Symbol -> REnv -> REnv
deleteREnv Symbol
x REnv
rE = REnv
rE REnv
-> (HashMap Symbol SpecType -> HashMap Symbol SpecType) -> REnv
`updREnvLocal` (Symbol -> HashMap Symbol SpecType -> HashMap Symbol SpecType
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> HashMap k v
M.delete Symbol
x)

insertREnv :: F.Symbol -> SpecType -> REnv -> REnv
insertREnv :: Symbol -> SpecType -> REnv -> REnv
insertREnv Symbol
x SpecType
y REnv
rE = {- trace ("insertREnv: " ++ show x) $ -} REnv
rE REnv
-> (HashMap Symbol SpecType -> HashMap Symbol SpecType) -> REnv
`updREnvLocal` (Symbol
-> SpecType -> HashMap Symbol SpecType -> HashMap Symbol SpecType
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
M.insert Symbol
x SpecType
y)

lookupREnv :: F.Symbol -> REnv -> Maybe SpecType
lookupREnv :: Symbol -> REnv -> Maybe SpecType
lookupREnv Symbol
x REnv
rE = [Maybe SpecType] -> Maybe SpecType
forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, MonadPlus m) =>
t (m a) -> m a
msum ([Maybe SpecType] -> Maybe SpecType)
-> [Maybe SpecType] -> Maybe SpecType
forall a b. (a -> b) -> a -> b
$ Symbol -> HashMap Symbol SpecType -> Maybe SpecType
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
M.lookup Symbol
x (HashMap Symbol SpecType -> Maybe SpecType)
-> [HashMap Symbol SpecType] -> [Maybe SpecType]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> REnv -> [HashMap Symbol SpecType]
renvMaps REnv
rE

memberREnv :: F.Symbol -> REnv -> Bool
memberREnv :: Symbol -> REnv -> Bool
memberREnv Symbol
x REnv
rE = [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
or   ([Bool] -> Bool) -> [Bool] -> Bool
forall a b. (a -> b) -> a -> b
$ Symbol -> HashMap Symbol SpecType -> Bool
forall k a. (Eq k, Hashable k) => k -> HashMap k a -> Bool
M.member Symbol
x (HashMap Symbol SpecType -> Bool)
-> [HashMap Symbol SpecType] -> [Bool]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> REnv -> [HashMap Symbol SpecType]
renvMaps REnv
rE

globalREnv :: REnv -> REnv
globalREnv :: REnv -> REnv
globalREnv (REnv HashMap Symbol SpecType
gM HashMap Symbol SpecType
lM) = HashMap Symbol SpecType -> HashMap Symbol SpecType -> REnv
forall t. HashMap Symbol t -> HashMap Symbol t -> AREnv t
REnv HashMap Symbol SpecType
gM' HashMap Symbol SpecType
forall k v. HashMap k v
M.empty
  where
    gM' :: HashMap Symbol SpecType
gM'  = (SpecType -> SpecType -> SpecType)
-> HashMap Symbol SpecType
-> HashMap Symbol SpecType
-> HashMap Symbol SpecType
forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> HashMap k v -> HashMap k v -> HashMap k v
M.unionWith (\SpecType
_ SpecType
t -> SpecType
t) HashMap Symbol SpecType
gM HashMap Symbol SpecType
lM

renvMaps :: REnv -> [M.HashMap F.Symbol SpecType]
renvMaps :: REnv -> [HashMap Symbol SpecType]
renvMaps REnv
rE = [REnv -> HashMap Symbol SpecType
forall t. AREnv t -> HashMap Symbol t
reLocal REnv
rE, REnv -> HashMap Symbol SpecType
forall t. AREnv t -> HashMap Symbol t
reGlobal REnv
rE]

--------------------------------------------------------------------------------
localBindsOfType :: RRType r  -> REnv -> [F.Symbol]
--------------------------------------------------------------------------------
localBindsOfType :: RRType r -> REnv -> [Symbol]
localBindsOfType RRType r
tx REnv
γ = (Symbol, SpecType) -> Symbol
forall a b. (a, b) -> a
fst ((Symbol, SpecType) -> Symbol) -> [(Symbol, SpecType)] -> [Symbol]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> REnv -> [(Symbol, SpecType)]
localsREnv ((SpecType -> Bool) -> REnv -> REnv
filterREnv ((RType RTyCon RTyVar () -> RType RTyCon RTyVar () -> Bool
forall a. Eq a => a -> a -> Bool
== RRType r -> RType RTyCon RTyVar ()
forall c tv r. RType c tv r -> RType c tv ()
toRSort RRType r
tx) (RType RTyCon RTyVar () -> Bool)
-> (SpecType -> RType RTyCon RTyVar ()) -> SpecType -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SpecType -> RType RTyCon RTyVar ()
forall c tv r. RType c tv r -> RType c tv ()
toRSort) REnv
γ)

-- RJ: REnv-Split-Bug?
localsREnv :: REnv -> [(F.Symbol, SpecType)]
localsREnv :: REnv -> [(Symbol, SpecType)]
localsREnv = HashMap Symbol SpecType -> [(Symbol, SpecType)]
forall k v. HashMap k v -> [(k, v)]
M.toList (HashMap Symbol SpecType -> [(Symbol, SpecType)])
-> (REnv -> HashMap Symbol SpecType)
-> REnv
-> [(Symbol, SpecType)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. REnv -> HashMap Symbol SpecType
forall t. AREnv t -> HashMap Symbol t
reLocal

globalsREnv :: REnv -> [(F.Symbol, SpecType)]
globalsREnv :: REnv -> [(Symbol, SpecType)]
globalsREnv = HashMap Symbol SpecType -> [(Symbol, SpecType)]
forall k v. HashMap k v -> [(k, v)]
M.toList (HashMap Symbol SpecType -> [(Symbol, SpecType)])
-> (REnv -> HashMap Symbol SpecType)
-> REnv
-> [(Symbol, SpecType)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. REnv -> HashMap Symbol SpecType
forall t. AREnv t -> HashMap Symbol t
reGlobal

toListREnv :: REnv -> [(F.Symbol, SpecType)]
toListREnv :: REnv -> [(Symbol, SpecType)]
toListREnv REnv
re = REnv -> [(Symbol, SpecType)]
globalsREnv REnv
re [(Symbol, SpecType)]
-> [(Symbol, SpecType)] -> [(Symbol, SpecType)]
forall a. [a] -> [a] -> [a]
++ REnv -> [(Symbol, SpecType)]
localsREnv REnv
re

--------------------------------------------------------------------------------
extendEnvWithVV :: CGEnv -> SpecType -> CG CGEnv
--------------------------------------------------------------------------------
extendEnvWithVV :: CGEnv -> SpecType -> CG CGEnv
extendEnvWithVV CGEnv
γ SpecType
t
  | Symbol -> Bool
F.isNontrivialVV Symbol
vv Bool -> Bool -> Bool
&& Bool -> Bool
not (Symbol
vv Symbol -> REnv -> Bool
`memberREnv` (CGEnv -> REnv
renv CGEnv
γ))
  = CGEnv
γ CGEnv -> (String, Symbol, SpecType) -> CG CGEnv
+= (String
"extVV", Symbol
vv, SpecType
t)
  | Bool
otherwise
  = CGEnv -> CG CGEnv
forall (m :: * -> *) a. Monad m => a -> m a
return CGEnv
γ
  where
    vv :: Symbol
vv = SpecType -> Symbol
forall r c tv. Reftable r => RType c tv r -> Symbol
rTypeValueVar SpecType
t

addBinders :: CGEnv -> F.Symbol -> [(F.Symbol, SpecType)] -> CG CGEnv
addBinders :: CGEnv -> Symbol -> [(Symbol, SpecType)] -> CG CGEnv
addBinders CGEnv
γ0 Symbol
x' [(Symbol, SpecType)]
cbs   = (CGEnv -> (String, Symbol, SpecType) -> CG CGEnv)
-> CGEnv -> [(String, Symbol, SpecType)] -> CG CGEnv
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM CGEnv -> (String, Symbol, SpecType) -> CG CGEnv
(++=) (CGEnv
γ0 CGEnv -> Symbol -> CGEnv
-= Symbol
x') [(String
"addBinders", Symbol
x, SpecType
t) | (Symbol
x, SpecType
t) <- [(Symbol, SpecType)]
cbs]

addBind :: SrcSpan -> F.Symbol -> F.SortedReft -> CG ((F.Symbol, F.Sort), F.BindId)
addBind :: SrcSpan -> Symbol -> SortedReft -> CG ((Symbol, Sort), BindId)
addBind SrcSpan
l Symbol
x SortedReft
r = do
  CGInfo
st          <- StateT CGInfo Identity CGInfo
forall s (m :: * -> *). MonadState s m => m s
get
  let (BindId
i, BindEnv
bs') = Symbol -> SortedReft -> BindEnv -> (BindId, BindEnv)
F.insertBindEnv Symbol
x SortedReft
r (CGInfo -> BindEnv
binds CGInfo
st)
  CGInfo -> StateT CGInfo Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put          (CGInfo -> StateT CGInfo Identity ())
-> CGInfo -> StateT CGInfo Identity ()
forall a b. (a -> b) -> a -> b
$ CGInfo
st { binds :: BindEnv
binds = BindEnv
bs' } { bindSpans :: HashMap BindId SrcSpan
bindSpans = BindId
-> SrcSpan -> HashMap BindId SrcSpan -> HashMap BindId SrcSpan
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
M.insert BindId
i SrcSpan
l (CGInfo -> HashMap BindId SrcSpan
bindSpans CGInfo
st) }
  ((Symbol, Sort), BindId) -> CG ((Symbol, Sort), BindId)
forall (m :: * -> *) a. Monad m => a -> m a
return ((Symbol
x, SortedReft -> Sort
F.sr_sort SortedReft
r), {- traceShow ("addBind: " ++ showpp x) -} BindId
i)

addClassBind :: CGEnv -> SrcSpan -> SpecType -> CG [((F.Symbol, F.Sort), F.BindId)]
addClassBind :: CGEnv -> SrcSpan -> SpecType -> CG [((Symbol, Sort), BindId)]
addClassBind CGEnv
γ SrcSpan
l = ((Symbol, SortedReft) -> CG ((Symbol, Sort), BindId))
-> [(Symbol, SortedReft)] -> CG [((Symbol, Sort), BindId)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM ((Symbol -> SortedReft -> CG ((Symbol, Sort), BindId))
-> (Symbol, SortedReft) -> CG ((Symbol, Sort), BindId)
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (SrcSpan -> Symbol -> SortedReft -> CG ((Symbol, Sort), BindId)
addBind SrcSpan
l)) ([(Symbol, SortedReft)] -> CG [((Symbol, Sort), BindId)])
-> (SpecType -> [(Symbol, SortedReft)])
-> SpecType
-> CG [((Symbol, Sort), BindId)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TCEmb TyCon -> SpecType -> [(Symbol, SortedReft)]
classBinds (CGEnv -> TCEmb TyCon
emb CGEnv
γ)

{- see tests/pos/polyfun for why you need everything in fixenv -}
addCGEnv :: (SpecType -> SpecType) -> CGEnv -> (String, F.Symbol, SpecType) -> CG CGEnv
addCGEnv :: (SpecType -> SpecType)
-> CGEnv -> (String, Symbol, SpecType) -> CG CGEnv
addCGEnv SpecType -> SpecType
tx CGEnv
γ (String
eMsg, Symbol
x, REx Symbol
y SpecType
tyy SpecType
tyx) = do
  Symbol
y' <- StateT CGInfo Identity Symbol
forall (m :: * -> *) a. Freshable m a => m a
fresh
  CGEnv
γ' <- (SpecType -> SpecType)
-> CGEnv -> (String, Symbol, SpecType) -> CG CGEnv
addCGEnv SpecType -> SpecType
tx CGEnv
γ (String
eMsg, Symbol
y', SpecType
tyy)
  (SpecType -> SpecType)
-> CGEnv -> (String, Symbol, SpecType) -> CG CGEnv
addCGEnv SpecType -> SpecType
tx CGEnv
γ' (String
eMsg, Symbol
x, SpecType
tyx SpecType -> (Symbol, Expr) -> SpecType
forall a. Subable a => a -> (Symbol, Expr) -> a
`F.subst1` (Symbol
y, Symbol -> Expr
F.EVar Symbol
y'))

addCGEnv SpecType -> SpecType
tx CGEnv
γ (String
eMsg, Symbol
x, RAllE Symbol
yy SpecType
tyy SpecType
tyx)
  = (SpecType -> SpecType)
-> CGEnv -> (String, Symbol, SpecType) -> CG CGEnv
addCGEnv SpecType -> SpecType
tx CGEnv
γ (String
eMsg, Symbol
x, SpecType
t)
  where
    xs :: [Symbol]
xs            = SpecType -> REnv -> [Symbol]
forall r. RRType r -> REnv -> [Symbol]
localBindsOfType SpecType
tyy (CGEnv -> REnv
renv CGEnv
γ)
    t :: SpecType
t             = (SpecType -> SpecType -> SpecType)
-> SpecType -> [SpecType] -> SpecType
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
L.foldl' SpecType -> SpecType -> SpecType
forall r. Reftable r => r -> r -> r
F.meet SpecType
ttrue [ SpecType
tyx' SpecType -> (Symbol, Expr) -> SpecType
forall a. Subable a => a -> (Symbol, Expr) -> a
`F.subst1` (Symbol
yy, Symbol -> Expr
F.EVar Symbol
x) | Symbol
x <- [Symbol]
xs]
    (SpecType
tyx', SpecType
ttrue) = Symbol -> SpecType -> (SpecType, SpecType)
splitXRelatedRefs Symbol
yy SpecType
tyx

addCGEnv SpecType -> SpecType
tx CGEnv
γ (String
_, Symbol
x, SpecType
t') = do
  Integer
idx   <- StateT CGInfo Identity Integer
forall (m :: * -> *) a. Freshable m a => m a
fresh
  -- allowHOBinders <- allowHO <$> get
  let t :: SpecType
t  = SpecType -> SpecType
tx (SpecType -> SpecType) -> SpecType -> SpecType
forall a b. (a -> b) -> a -> b
$ Integer -> SpecType -> SpecType
normalize Integer
idx SpecType
t'
  let l :: SrcSpan
l  = CGEnv -> SrcSpan
getLocation CGEnv
γ
  let γ' :: CGEnv
γ' = CGEnv
γ { renv :: REnv
renv = Symbol -> SpecType -> REnv -> REnv
insertREnv Symbol
x SpecType
t (CGEnv -> REnv
renv CGEnv
γ) }
  Templates
tem   <- CG Templates
getTemplates 
  [((Symbol, Sort), BindId)]
is    <- (:) (((Symbol, Sort), BindId)
 -> [((Symbol, Sort), BindId)] -> [((Symbol, Sort), BindId)])
-> CG ((Symbol, Sort), BindId)
-> StateT
     CGInfo
     Identity
     ([((Symbol, Sort), BindId)] -> [((Symbol, Sort), BindId)])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SrcSpan -> Symbol -> SortedReft -> CG ((Symbol, Sort), BindId)
addBind SrcSpan
l Symbol
x (CGEnv -> Templates -> SpecType -> SortedReft
forall r.
(PPrint r, Reftable r, SubsTy RTyVar (RType RTyCon RTyVar ()) r,
 Reftable (RTProp RTyCon RTyVar r)) =>
CGEnv -> Templates -> RRType r -> SortedReft
rTypeSortedReft' CGEnv
γ' Templates
tem SpecType
t) StateT
  CGInfo
  Identity
  ([((Symbol, Sort), BindId)] -> [((Symbol, Sort), BindId)])
-> CG [((Symbol, Sort), BindId)] -> CG [((Symbol, Sort), BindId)]
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> CGEnv -> SrcSpan -> SpecType -> CG [((Symbol, Sort), BindId)]
addClassBind CGEnv
γ' SrcSpan
l SpecType
t
  CGEnv -> CG CGEnv
forall (m :: * -> *) a. Monad m => a -> m a
return (CGEnv -> CG CGEnv) -> CGEnv -> CG CGEnv
forall a b. (a -> b) -> a -> b
$ CGEnv
γ' { fenv :: FEnv
fenv = FEnv -> [((Symbol, Sort), BindId)] -> FEnv
insertsFEnv (CGEnv -> FEnv
fenv CGEnv
γ) [((Symbol, Sort), BindId)]
is }

rTypeSortedReft' :: (PPrint r, F.Reftable r, SubsTy RTyVar RSort r, F.Reftable (RTProp RTyCon RTyVar r))
    => CGEnv -> F.Templates -> RRType r -> F.SortedReft
rTypeSortedReft' :: CGEnv -> Templates -> RRType r -> SortedReft
rTypeSortedReft' CGEnv
γ Templates
t 
  = SEnv Sort -> Templates -> SortedReft -> SortedReft
pruneUnsortedReft (FEnv -> SEnv Sort
feEnv (FEnv -> SEnv Sort) -> FEnv -> SEnv Sort
forall a b. (a -> b) -> a -> b
$ CGEnv -> FEnv
fenv CGEnv
γ) Templates
t (SortedReft -> SortedReft)
-> (RRType r -> SortedReft) -> RRType r -> SortedReft
forall b c a. (b -> c) -> (a -> b) -> a -> c
. RRType r -> SortedReft
forall r.
(PPrint r, SubsTy RTyVar (RType RTyCon RTyVar ()) r, Reftable r,
 Reftable (RTProp RTyCon RTyVar r)) =>
RRType r -> SortedReft
f
   where
   f :: RRType r -> SortedReft
f         = TCEmb TyCon -> RRType r -> SortedReft
forall r.
(PPrint r, Reftable r, SubsTy RTyVar (RType RTyCon RTyVar ()) r,
 Reftable (RTProp RTyCon RTyVar r)) =>
TCEmb TyCon -> RRType r -> SortedReft
rTypeSortedReft (CGEnv -> TCEmb TyCon
emb CGEnv
γ)


normalize :: Integer -> SpecType -> SpecType
normalize :: Integer -> SpecType -> SpecType
normalize Integer
idx = Integer -> SpecType -> SpecType
normalizeVV Integer
idx (SpecType -> SpecType)
-> (SpecType -> SpecType) -> SpecType -> SpecType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SpecType -> SpecType
forall c tv r. OkRT c tv r => RType c tv r -> RType c tv r
normalizePds

normalizeVV :: Integer -> SpecType -> SpecType
normalizeVV :: Integer -> SpecType -> SpecType
normalizeVV Integer
idx t :: SpecType
t@(RApp RTyCon
_ [SpecType]
_ [RTProp RTyCon RTyVar RReft]
_ RReft
_)
  | Bool -> Bool
not (Symbol -> Bool
F.isNontrivialVV (SpecType -> Symbol
forall r c tv. Reftable r => RType c tv r -> Symbol
rTypeValueVar SpecType
t))
  = SpecType -> Symbol -> SpecType
forall c (f :: * -> *) tv.
(TyConable c, Reftable (f Reft), Functor f) =>
RType c tv (f Reft) -> Symbol -> RType c tv (f Reft)
shiftVV SpecType
t (Maybe Integer -> Symbol
F.vv (Maybe Integer -> Symbol) -> Maybe Integer -> Symbol
forall a b. (a -> b) -> a -> b
$ Integer -> Maybe Integer
forall a. a -> Maybe a
Just Integer
idx)

normalizeVV Integer
_ SpecType
t
  = SpecType
t

--------------------------------------------------------------------------------
(+=) :: CGEnv -> (String, F.Symbol, SpecType) -> CG CGEnv
--------------------------------------------------------------------------------
CGEnv
γ += :: CGEnv -> (String, Symbol, SpecType) -> CG CGEnv
+= (String
eMsg, Symbol
x, SpecType
r)
  | Symbol
x Symbol -> Symbol -> Bool
forall a. Eq a => a -> a -> Bool
== Symbol
F.dummySymbol
  = CGEnv -> CG CGEnv
forall (m :: * -> *) a. Monad m => a -> m a
return CGEnv
γ
  -- // | x `memberREnv` (renv γ)
  -- // = _dupBindErr x γ
  | Bool
otherwise
  =  CGEnv
γ CGEnv -> (String, Symbol, SpecType) -> CG CGEnv
++= (String
eMsg, Symbol
x, SpecType
r)

_dupBindError :: String -> F.Symbol -> CGEnv -> SpecType -> a
_dupBindError :: String -> Symbol -> CGEnv -> SpecType -> a
_dupBindError String
eMsg Symbol
x CGEnv
γ SpecType
r = Maybe SrcSpan -> String -> a
forall a. Maybe SrcSpan -> String -> a
panic Maybe SrcSpan
forall a. Maybe a
Nothing String
s
  where
    s :: String
s = [String] -> String
unlines [ String
eMsg String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" Duplicate binding for " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Symbol -> String
F.symbolString Symbol
x
                , String
"   New: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ SpecType -> String
forall a. PPrint a => a -> String
showpp SpecType
r
                , String
"   Old: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ Maybe SpecType -> String
forall a. PPrint a => a -> String
showpp (Symbol
x Symbol -> REnv -> Maybe SpecType
`lookupREnv` (CGEnv -> REnv
renv CGEnv
γ)) ]

--------------------------------------------------------------------------------
globalize :: CGEnv -> CGEnv
--------------------------------------------------------------------------------
globalize :: CGEnv -> CGEnv
globalize CGEnv
γ = CGEnv
γ {renv :: REnv
renv = REnv -> REnv
globalREnv (CGEnv -> REnv
renv CGEnv
γ)}

--------------------------------------------------------------------------------
(++=) :: CGEnv -> (String, F.Symbol, SpecType) -> CG CGEnv
--------------------------------------------------------------------------------
++= :: CGEnv -> (String, Symbol, SpecType) -> CG CGEnv
(++=) CGEnv
γ (String
eMsg, Symbol
x, SpecType
t)
  = (SpecType -> SpecType)
-> CGEnv -> (String, Symbol, SpecType) -> CG CGEnv
addCGEnv (RTyConInv -> SpecType -> SpecType
addRTyConInv (([RInv] -> [RInv] -> [RInv]) -> RTyConInv -> RTyConInv -> RTyConInv
forall k v.
(Eq k, Hashable k) =>
(v -> v -> v) -> HashMap k v -> HashMap k v -> HashMap k v
M.unionWith [RInv] -> [RInv] -> [RInv]
forall a. Monoid a => a -> a -> a
mappend (CGEnv -> RTyConInv
invs CGEnv
γ) (CGEnv -> RTyConInv
ial CGEnv
γ))) CGEnv
γ (String
eMsg, Symbol
x, SpecType
t)

--------------------------------------------------------------------------------
addSEnv :: CGEnv -> (String, F.Symbol, SpecType) -> CG CGEnv
--------------------------------------------------------------------------------
addSEnv :: CGEnv -> (String, Symbol, SpecType) -> CG CGEnv
addSEnv CGEnv
γ = (SpecType -> SpecType)
-> CGEnv -> (String, Symbol, SpecType) -> CG CGEnv
addCGEnv (RTyConInv -> SpecType -> SpecType
addRTyConInv (CGEnv -> RTyConInv
invs CGEnv
γ)) CGEnv
γ

addEEnv :: CGEnv -> (F.Symbol, SpecType) -> CG CGEnv
addEEnv :: CGEnv -> (Symbol, SpecType) -> CG CGEnv
addEEnv CGEnv
γ (Symbol
x,SpecType
t')= do
  Integer
idx   <- StateT CGInfo Identity Integer
forall (m :: * -> *) a. Freshable m a => m a
fresh
  -- allowHOBinders <- allowHO <$> get
  let t :: SpecType
t  = RTyConInv -> SpecType -> SpecType
addRTyConInv (CGEnv -> RTyConInv
invs CGEnv
γ) (SpecType -> SpecType) -> SpecType -> SpecType
forall a b. (a -> b) -> a -> b
$ Integer -> SpecType -> SpecType
normalize Integer
idx SpecType
t'
  let l :: SrcSpan
l  = CGEnv -> SrcSpan
getLocation CGEnv
γ
  let γ' :: CGEnv
γ' = CGEnv
γ { renv :: REnv
renv = Symbol -> SpecType -> REnv -> REnv
insertREnv Symbol
x SpecType
t (CGEnv -> REnv
renv CGEnv
γ) }
  Templates
tem   <- CG Templates
getTemplates
  [((Symbol, Sort), BindId)]
is    <- (:) (((Symbol, Sort), BindId)
 -> [((Symbol, Sort), BindId)] -> [((Symbol, Sort), BindId)])
-> CG ((Symbol, Sort), BindId)
-> StateT
     CGInfo
     Identity
     ([((Symbol, Sort), BindId)] -> [((Symbol, Sort), BindId)])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SrcSpan -> Symbol -> SortedReft -> CG ((Symbol, Sort), BindId)
addBind SrcSpan
l Symbol
x (CGEnv -> Templates -> SpecType -> SortedReft
forall r.
(PPrint r, Reftable r, SubsTy RTyVar (RType RTyCon RTyVar ()) r,
 Reftable (RTProp RTyCon RTyVar r)) =>
CGEnv -> Templates -> RRType r -> SortedReft
rTypeSortedReft' CGEnv
γ' Templates
tem SpecType
t) StateT
  CGInfo
  Identity
  ([((Symbol, Sort), BindId)] -> [((Symbol, Sort), BindId)])
-> CG [((Symbol, Sort), BindId)] -> CG [((Symbol, Sort), BindId)]
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> CGEnv -> SrcSpan -> SpecType -> CG [((Symbol, Sort), BindId)]
addClassBind CGEnv
γ' SrcSpan
l SpecType
t
  (CGInfo -> CGInfo) -> StateT CGInfo Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify (\CGInfo
s -> CGInfo
s { ebinds :: [BindId]
ebinds = CGInfo -> [BindId]
ebinds CGInfo
s [BindId] -> [BindId] -> [BindId]
forall a. [a] -> [a] -> [a]
++ (((Symbol, Sort), BindId) -> BindId
forall a b. (a, b) -> b
snd (((Symbol, Sort), BindId) -> BindId)
-> [((Symbol, Sort), BindId)] -> [BindId]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [((Symbol, Sort), BindId)]
is)})
  CGEnv -> CG CGEnv
forall (m :: * -> *) a. Monad m => a -> m a
return (CGEnv -> CG CGEnv) -> CGEnv -> CG CGEnv
forall a b. (a -> b) -> a -> b
$ CGEnv
γ' { fenv :: FEnv
fenv = FEnv -> [((Symbol, Sort), BindId)] -> FEnv
insertsFEnv (CGEnv -> FEnv
fenv CGEnv
γ) [((Symbol, Sort), BindId)]
is }


(+++=) :: (CGEnv, String) -> (F.Symbol, CoreExpr, SpecType) -> CG CGEnv
(CGEnv
γ, String
_) +++= :: (CGEnv, String) -> (Symbol, CoreExpr, SpecType) -> CG CGEnv
+++= (Symbol
x, CoreExpr
e, SpecType
t) = (CGEnv
γ {lcb :: HashMap Symbol CoreExpr
lcb = Symbol
-> CoreExpr -> HashMap Symbol CoreExpr -> HashMap Symbol CoreExpr
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
M.insert Symbol
x CoreExpr
e (CGEnv -> HashMap Symbol CoreExpr
lcb CGEnv
γ) }) CGEnv -> (String, Symbol, SpecType) -> CG CGEnv
+= (String
"+++=", Symbol
x, SpecType
t)

(-=) :: CGEnv -> F.Symbol -> CGEnv
CGEnv
γ -= :: CGEnv -> Symbol -> CGEnv
-= Symbol
x =  CGEnv
γ { renv :: REnv
renv = Symbol -> REnv -> REnv
deleteREnv Symbol
x (CGEnv -> REnv
renv CGEnv
γ)
            , lcb :: HashMap Symbol CoreExpr
lcb  = Symbol -> HashMap Symbol CoreExpr -> HashMap Symbol CoreExpr
forall k v. (Eq k, Hashable k) => k -> HashMap k v -> HashMap k v
M.delete   Symbol
x (CGEnv -> HashMap Symbol CoreExpr
lcb  CGEnv
γ)
            -- , fenv = removeFEnv x (fenv γ)
            }

(?=) :: (?callStack :: CallStack) => CGEnv -> F.Symbol -> Maybe SpecType
CGEnv
γ ?= :: CGEnv -> Symbol -> Maybe SpecType
?= Symbol
x  = Symbol -> REnv -> Maybe SpecType
lookupREnv Symbol
x (CGEnv -> REnv
renv CGEnv
γ)

------------------------------------------------------------------------
setLocation :: CGEnv -> Sp.Span -> CGEnv
------------------------------------------------------------------------
setLocation :: CGEnv -> Span -> CGEnv
setLocation CGEnv
γ Span
p = CGEnv
γ { cgLoc :: SpanStack
cgLoc = Span -> SpanStack -> SpanStack
Sp.push Span
p (SpanStack -> SpanStack) -> SpanStack -> SpanStack
forall a b. (a -> b) -> a -> b
$ CGEnv -> SpanStack
cgLoc CGEnv
γ }

------------------------------------------------------------------------
setBind :: CGEnv -> Var -> CGEnv
------------------------------------------------------------------------
setBind :: CGEnv -> Var -> CGEnv
setBind CGEnv
γ Var
x = CGEnv
γ CGEnv -> Span -> CGEnv
`setLocation` Var -> Span
Sp.Var Var
x CGEnv -> Var -> CGEnv
`setBind'` Var
x

setBind' :: CGEnv -> Tg.TagKey -> CGEnv
setBind' :: CGEnv -> Var -> CGEnv
setBind' CGEnv
γ Var
k
  | Var -> TagEnv -> Bool
Tg.memTagEnv Var
k (CGEnv -> TagEnv
tgEnv CGEnv
γ) = CGEnv
γ { tgKey :: Maybe Var
tgKey = Var -> Maybe Var
forall a. a -> Maybe a
Just Var
k }
  | Bool
otherwise                = CGEnv
γ

------------------------------------------------------------------------
setRecs :: CGEnv -> [Var] -> CGEnv
------------------------------------------------------------------------
setRecs :: CGEnv -> [Var] -> CGEnv
setRecs CGEnv
γ [Var]
xs   = CGEnv
γ { recs :: HashSet Var
recs = (HashSet Var -> Var -> HashSet Var)
-> HashSet Var -> [Var] -> HashSet Var
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
L.foldl' ((Var -> HashSet Var -> HashSet Var)
-> HashSet Var -> Var -> HashSet Var
forall a b c. (a -> b -> c) -> b -> a -> c
flip Var -> HashSet Var -> HashSet Var
forall a. (Eq a, Hashable a) => a -> HashSet a -> HashSet a
S.insert) (CGEnv -> HashSet Var
recs CGEnv
γ) [Var]
xs }

------------------------------------------------------------------------
setTRec :: CGEnv -> [(Var, SpecType)] -> CGEnv
------------------------------------------------------------------------
setTRec :: CGEnv -> [(Var, SpecType)] -> CGEnv
setTRec CGEnv
γ [(Var, SpecType)]
xts  = CGEnv
γ' {trec :: Maybe (HashMap Symbol SpecType)
trec = HashMap Symbol SpecType -> Maybe (HashMap Symbol SpecType)
forall a. a -> Maybe a
Just (HashMap Symbol SpecType -> Maybe (HashMap Symbol SpecType))
-> HashMap Symbol SpecType -> Maybe (HashMap Symbol SpecType)
forall a b. (a -> b) -> a -> b
$ [(Symbol, SpecType)] -> HashMap Symbol SpecType
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
M.fromList [(Symbol, SpecType)]
xts' HashMap Symbol SpecType
-> HashMap Symbol SpecType -> HashMap Symbol SpecType
forall k v.
(Eq k, Hashable k) =>
HashMap k v -> HashMap k v -> HashMap k v
`M.union` HashMap Symbol SpecType
trec'}
  where
    γ' :: CGEnv
γ'         = CGEnv
γ CGEnv -> [Var] -> CGEnv
`setRecs` ((Var, SpecType) -> Var
forall a b. (a, b) -> a
fst ((Var, SpecType) -> Var) -> [(Var, SpecType)] -> [Var]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Var, SpecType)]
xts)
    trec' :: HashMap Symbol SpecType
trec'      = HashMap Symbol SpecType
-> Maybe (HashMap Symbol SpecType) -> HashMap Symbol SpecType
forall a. a -> Maybe a -> a
fromMaybe HashMap Symbol SpecType
forall k v. HashMap k v
M.empty (Maybe (HashMap Symbol SpecType) -> HashMap Symbol SpecType)
-> Maybe (HashMap Symbol SpecType) -> HashMap Symbol SpecType
forall a b. (a -> b) -> a -> b
$ CGEnv -> Maybe (HashMap Symbol SpecType)
trec CGEnv
γ
    xts' :: [(Symbol, SpecType)]
xts'       = (Var -> Symbol) -> (Var, SpecType) -> (Symbol, SpecType)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first Var -> Symbol
forall a. Symbolic a => a -> Symbol
F.symbol ((Var, SpecType) -> (Symbol, SpecType))
-> [(Var, SpecType)] -> [(Symbol, SpecType)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Var, SpecType)]
xts

------------------------------------------------------------------------
getLocation :: CGEnv -> SrcSpan
------------------------------------------------------------------------
getLocation :: CGEnv -> SrcSpan
getLocation = SpanStack -> SrcSpan
Sp.srcSpan (SpanStack -> SrcSpan) -> (CGEnv -> SpanStack) -> CGEnv -> SrcSpan
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CGEnv -> SpanStack
cgLoc