{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE UndecidableInstances  #-}
{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE ConstraintKinds       #-}

{-# OPTIONS_GHC -Wno-orphans #-}

module Language.Haskell.Liquid.Constraint.Fresh
  ( -- module Language.Haskell.Liquid.Types.Fresh
    -- , 
    refreshArgsTop
  , freshTyType
  , freshTyExpr
  , trueTy
  , addKuts
  )
  where

-- import           Data.Maybe                    (catMaybes) -- , fromJust, isJust)
-- import           Data.Bifunctor
-- import qualified Data.List                      as L
import qualified Data.HashMap.Strict            as M
import qualified Data.HashSet                   as S
import           Data.Hashable
import           Control.Monad.State            (gets, get, put, modify)
import           Control.Monad                  (when, (>=>))
import           Prelude                        hiding (error)

import qualified Language.Fixpoint.Types as F
import           Language.Fixpoint.Types.Visitor (kvarsExpr)
import           Language.Haskell.Liquid.Types.Fresh
import           Language.Haskell.Liquid.Types.RefType
import           Language.Haskell.Liquid.Types.RType
import           Language.Haskell.Liquid.Types.RTypeOp
import           Language.Haskell.Liquid.Types.Types
import           Language.Haskell.Liquid.Constraint.Types
import qualified Language.Haskell.Liquid.GHC.Misc as GM
import           Language.Haskell.Liquid.UX.Config
import           Liquid.GHC.API as Ghc hiding (get)

--------------------------------------------------------------------------------
-- | This is all hardwiring stuff to CG ----------------------------------------
--------------------------------------------------------------------------------
instance Freshable CG Integer where
  fresh :: CG Integer
fresh = do s <- StateT CGInfo Identity CGInfo
forall s (m :: * -> *). MonadState s m => m s
get
             let n = CGInfo -> Integer
freshIndex CGInfo
s
             put $ s { freshIndex = n + 1 }
             return n

--------------------------------------------------------------------------------
refreshArgsTop :: (Var, SpecType) -> CG SpecType
--------------------------------------------------------------------------------
refreshArgsTop :: (Var, SpecType) -> CG SpecType
refreshArgsTop (Var
x, SpecType
t)
  = do (t', su) <- SpecType -> StateT CGInfo Identity (SpecType, Subst)
forall (m :: * -> *). FreshM m => SpecType -> m (SpecType, Subst)
refreshArgsSub SpecType
t
       modify $ \CGInfo
s -> CGInfo
s {termExprs = M.adjust (F.subst su <$>) x $ termExprs s}
       return t'

--------------------------------------------------------------------------------
-- | Generation: Freshness -----------------------------------------------------
--------------------------------------------------------------------------------

-- | Right now, we generate NO new pvars. Rather than clutter code
--   with `uRType` calls, put it in one place where the above
--   invariant is /obviously/ enforced.
--   Constraint generation should ONLY use @freshTyType@ and @freshTyExpr@

freshTyType        :: Bool -> KVKind -> CoreExpr -> Type -> CG SpecType
freshTyType :: Bool -> KVKind -> CoreExpr -> Type -> CG SpecType
freshTyType Bool
allowTC KVKind
k CoreExpr
e Type
τ  =  [Char] -> SpecType -> SpecType
forall a. PPrint a => [Char] -> a -> a
F.notracepp ([Char]
"freshTyType: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ KVKind -> [Char]
forall a. PPrint a => a -> [Char]
F.showpp KVKind
k [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ CoreExpr -> [Char]
forall a. Outputable a => a -> [Char]
GM.showPpr CoreExpr
e)
                   (SpecType -> SpecType) -> CG SpecType -> CG SpecType
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Bool -> KVKind -> SpecType -> CG SpecType
freshTyReftype Bool
allowTC KVKind
k (Type -> SpecType
forall r. Monoid r => Type -> RRType r
ofType Type
τ)

freshTyExpr        :: Bool -> KVKind -> CoreExpr -> Type -> CG SpecType
freshTyExpr :: Bool -> KVKind -> CoreExpr -> Type -> CG SpecType
freshTyExpr Bool
allowTC KVKind
k CoreExpr
e Type
_  = Bool -> KVKind -> SpecType -> CG SpecType
freshTyReftype Bool
allowTC KVKind
k (SpecType -> CG SpecType) -> SpecType -> CG SpecType
forall a b. (a -> b) -> a -> b
$ CoreExpr -> SpecType
exprRefType CoreExpr
e

freshTyReftype     :: Bool -> KVKind -> SpecType -> CG SpecType
freshTyReftype :: Bool -> KVKind -> SpecType -> CG SpecType
freshTyReftype Bool
allowTC KVKind
k SpecType
t = do
    st <- SpecType -> CG SpecType
fixTy SpecType
t CG SpecType -> (SpecType -> CG SpecType) -> CG SpecType
forall a b.
StateT CGInfo Identity a
-> (a -> StateT CGInfo Identity b) -> StateT CGInfo Identity b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Bool -> SpecType -> CG SpecType
forall (m :: * -> *) a. Freshable m a => Bool -> a -> m a
refresh Bool
allowTC
    addKVars k st
    return st

-- | Used to generate "cut" kvars for fixpoint. Typically, KVars for recursive
--   definitions, and also to update the KVar profile.
addKVars        :: KVKind -> SpecType -> CG ()
addKVars :: KVKind -> SpecType -> StateT CGInfo Identity ()
addKVars !KVKind
k !SpecType
t  = do
    cfg <- (CGInfo -> Config) -> StateT CGInfo Identity Config
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (TargetInfo -> Config
forall t. HasConfig t => t -> Config
getConfig (TargetInfo -> Config)
-> (CGInfo -> TargetInfo) -> CGInfo -> Config
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CGInfo -> TargetInfo
ghcI)
    when True          $ modify $ \CGInfo
s -> CGInfo
s { kvProf = updKVProf k ks (kvProf s) }
    when (isKut cfg k) $ addKuts k t
  where
    ks :: Kuts
ks         = HashSet KVar -> Kuts
F.KS (HashSet KVar -> Kuts) -> HashSet KVar -> Kuts
forall a b. (a -> b) -> a -> b
$ [KVar] -> HashSet KVar
forall a. (Eq a, Hashable a) => [a] -> HashSet a
S.fromList ([KVar] -> HashSet KVar) -> [KVar] -> HashSet KVar
forall a b. (a -> b) -> a -> b
$ SpecType -> [KVar]
specTypeKVars SpecType
t

isKut :: Config -> KVKind -> Bool
isKut :: Config -> KVKind -> Bool
isKut Config
_  (RecBindE Var
_) = Bool
True
isKut Config
cfg KVKind
ProjectE    = Bool -> Bool
not (Config -> Bool
forall t. HasConfig t => t -> Bool
higherOrderFlag Config
cfg) -- see ISSUE 1034, tests/pos/T1034.hs
isKut Config
_    KVKind
_          = Bool
False

addKuts :: (PPrint a) => a -> SpecType -> CG ()
addKuts :: forall a. PPrint a => a -> SpecType -> StateT CGInfo Identity ()
addKuts a
_x SpecType
t = (CGInfo -> CGInfo) -> StateT CGInfo Identity ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify ((CGInfo -> CGInfo) -> StateT CGInfo Identity ())
-> (CGInfo -> CGInfo) -> StateT CGInfo Identity ()
forall a b. (a -> b) -> a -> b
$ \CGInfo
s -> CGInfo
s { kuts = mappend (F.KS ks) (kuts s)   }
  where
     ks' :: HashSet KVar
ks'     = [KVar] -> HashSet KVar
forall a. (Eq a, Hashable a) => [a] -> HashSet a
S.fromList ([KVar] -> HashSet KVar) -> [KVar] -> HashSet KVar
forall a b. (a -> b) -> a -> b
$ SpecType -> [KVar]
specTypeKVars SpecType
t
     ks :: HashSet KVar
ks
       | HashSet KVar -> Bool
forall a. HashSet a -> Bool
S.null HashSet KVar
ks' = HashSet KVar
ks'
       | Bool
otherwise  = {- F.tracepp ("addKuts: " ++ showpp _x) -} HashSet KVar
ks'

specTypeKVars :: SpecType -> [F.KVar]
specTypeKVars :: SpecType -> [KVar]
specTypeKVars = Bool
-> (SEnv SpecType -> RReft -> [KVar] -> [KVar])
-> [KVar]
-> SpecType
-> [KVar]
forall r c tv a.
(Reftable r, TyConable c) =>
Bool
-> (SEnv (RType c tv r) -> r -> a -> a) -> a -> RType c tv r -> a
foldReft Bool
False (\ SEnv SpecType
_ RReft
r [KVar]
ks -> Expr -> [KVar]
forall v. ExprV v -> [KVar]
kvarsExpr (ReftV Symbol -> Expr
forall v. ReftV v -> ExprV v
F.reftPred (ReftV Symbol -> Expr) -> ReftV Symbol -> Expr
forall a b. (a -> b) -> a -> b
$ RReft -> ReftV Symbol
forall v r. UReftV v r -> r
ur_reft RReft
r) [KVar] -> [KVar] -> [KVar]
forall a. [a] -> [a] -> [a]
++ [KVar]
ks) []

--------------------------------------------------------------------------------
trueTy  :: Bool -> Type -> CG SpecType
--------------------------------------------------------------------------------
trueTy :: Bool -> Type -> CG SpecType
trueTy Bool
allowTC = Type -> CG SpecType
ofType' (Type -> CG SpecType)
-> (SpecType -> CG SpecType) -> Type -> CG SpecType
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> Bool -> SpecType -> CG SpecType
forall (m :: * -> *) a. Freshable m a => Bool -> a -> m a
true Bool
allowTC

ofType' :: Type -> CG SpecType
ofType' :: Type -> CG SpecType
ofType' = SpecType -> CG SpecType
fixTy (SpecType -> CG SpecType)
-> (Type -> SpecType) -> Type -> CG SpecType
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Type -> SpecType
forall r. Monoid r => Type -> RRType r
ofType

fixTy :: SpecType -> CG SpecType
fixTy :: SpecType -> CG SpecType
fixTy SpecType
t = do tyi   <- (CGInfo -> TyConMap) -> StateT CGInfo Identity TyConMap
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets CGInfo -> TyConMap
tyConInfo
             tce   <- gets tyConEmbed
             return $ addTyConInfo tce tyi t

exprRefType :: CoreExpr -> SpecType
exprRefType :: CoreExpr -> SpecType
exprRefType = HashMap Var SpecType -> CoreExpr -> SpecType
exprRefType_ HashMap Var SpecType
forall k v. HashMap k v
M.empty

exprRefType_ :: M.HashMap Var SpecType -> CoreExpr -> SpecType
exprRefType_ :: HashMap Var SpecType -> CoreExpr -> SpecType
exprRefType_ HashMap Var SpecType
γ (Let Bind Var
b CoreExpr
e)
  = HashMap Var SpecType -> CoreExpr -> SpecType
exprRefType_ (HashMap Var SpecType -> Bind Var -> HashMap Var SpecType
bindRefType_ HashMap Var SpecType
γ Bind Var
b) CoreExpr
e

exprRefType_ HashMap Var SpecType
γ (Lam Var
α CoreExpr
e) | Var -> Bool
isTyVar Var
α
  = RTVUV Symbol RTyCon RTyVar -> SpecType -> RReft -> SpecType
forall v c tv r.
RTVUV v c tv -> RTypeV v c tv r -> r -> RTypeV v c tv r
RAllT (RTyVar -> RTVUV Symbol RTyCon RTyVar
forall tv s. tv -> RTVar tv s
makeRTVar (RTyVar -> RTVUV Symbol RTyCon RTyVar)
-> RTyVar -> RTVUV Symbol RTyCon RTyVar
forall a b. (a -> b) -> a -> b
$ Var -> RTyVar
rTyVar Var
α) (HashMap Var SpecType -> CoreExpr -> SpecType
exprRefType_ HashMap Var SpecType
γ CoreExpr
e) RReft
forall a. Monoid a => a
mempty

exprRefType_ HashMap Var SpecType
γ (Lam Var
x CoreExpr
e)
  = Symbol -> SpecType -> SpecType -> SpecType
forall r v c tv.
Monoid r =>
Symbol -> RTypeV v c tv r -> RTypeV v c tv r -> RTypeV v c tv r
rFun (Var -> Symbol
forall a. Symbolic a => a -> Symbol
F.symbol Var
x) (Type -> SpecType
forall r. Monoid r => Type -> RRType r
ofType (Type -> SpecType) -> Type -> SpecType
forall a b. (a -> b) -> a -> b
$ Var -> Type
varType Var
x) (HashMap Var SpecType -> CoreExpr -> SpecType
exprRefType_ HashMap Var SpecType
γ CoreExpr
e)

exprRefType_ HashMap Var SpecType
γ (Tick CoreTickish
_ CoreExpr
e)
  = HashMap Var SpecType -> CoreExpr -> SpecType
exprRefType_ HashMap Var SpecType
γ CoreExpr
e

exprRefType_ HashMap Var SpecType
γ (Var Var
x)
  = SpecType -> Var -> HashMap Var SpecType -> SpecType
forall k v. (Eq k, Hashable k) => v -> k -> HashMap k v -> v
M.lookupDefault (Type -> SpecType
forall r. Monoid r => Type -> RRType r
ofType (Type -> SpecType) -> Type -> SpecType
forall a b. (a -> b) -> a -> b
$ Var -> Type
varType Var
x) Var
x HashMap Var SpecType
γ

exprRefType_ HashMap Var SpecType
_ CoreExpr
e
  = Type -> SpecType
forall r. Monoid r => Type -> RRType r
ofType (Type -> SpecType) -> Type -> SpecType
forall a b. (a -> b) -> a -> b
$ HasDebugCallStack => CoreExpr -> Type
CoreExpr -> Type
exprType CoreExpr
e

bindRefType_ :: M.HashMap Var SpecType -> Bind Var -> M.HashMap Var SpecType
bindRefType_ :: HashMap Var SpecType -> Bind Var -> HashMap Var SpecType
bindRefType_ HashMap Var SpecType
γ (Rec [(Var, CoreExpr)]
xes)
  = HashMap Var SpecType -> [(Var, SpecType)] -> HashMap Var SpecType
forall k (t :: * -> *) v.
(Eq k, Foldable t, Hashable k) =>
HashMap k v -> t (k, v) -> HashMap k v
extendγ HashMap Var SpecType
γ [(Var
x, HashMap Var SpecType -> CoreExpr -> SpecType
exprRefType_ HashMap Var SpecType
γ CoreExpr
e) | (Var
x,CoreExpr
e) <- [(Var, CoreExpr)]
xes]

bindRefType_ HashMap Var SpecType
γ (NonRec Var
x CoreExpr
e)
  = HashMap Var SpecType -> [(Var, SpecType)] -> HashMap Var SpecType
forall k (t :: * -> *) v.
(Eq k, Foldable t, Hashable k) =>
HashMap k v -> t (k, v) -> HashMap k v
extendγ HashMap Var SpecType
γ [(Var
x, HashMap Var SpecType -> CoreExpr -> SpecType
exprRefType_ HashMap Var SpecType
γ CoreExpr
e)]

extendγ :: (Eq k, Foldable t, Hashable k)
        => M.HashMap k v
        -> t (k, v)
        -> M.HashMap k v
extendγ :: forall k (t :: * -> *) v.
(Eq k, Foldable t, Hashable k) =>
HashMap k v -> t (k, v) -> HashMap k v
extendγ HashMap k v
γ t (k, v)
xts
  = ((k, v) -> HashMap k v -> HashMap k v)
-> HashMap k v -> t (k, v) -> HashMap k v
forall a b. (a -> b -> b) -> b -> t a -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\(k
x,v
t) HashMap k v
m -> k -> v -> HashMap k v -> HashMap k v
forall k v.
(Eq k, Hashable k) =>
k -> v -> HashMap k v -> HashMap k v
M.insert k
x v
t HashMap k v
m) HashMap k v
γ t (k, v)
xts