Expand.hs

-- | This module has the code for applying refinement (and) type aliases
--   and the pipeline for "cooking" a @BareType@ into a @SpecType@.
--   TODO: _only_ export `makeRTEnv`, `cookSpecType` and maybe `qualifyExpand`...

{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# OPTIONS_GHC -Wno-x-partial #-}

module Language.Haskell.Liquid.Bare.Expand
  ( -- * Create alias expansion environment
    makeRTEnv

    -- * Expand
  , Expand(expand)

    -- * Converting BareType to SpecType
  , cookSpecType
  , cookSpecTypeE
  , specExpandType

    -- * Re-exported for data-constructors
  , plugHoles
  ) where

import Prelude hiding (error)
import Data.Graph hiding (Graph)
import Data.Maybe

import           Control.Monad
import           Control.Monad.Identity
import           Control.Monad.State
import           Data.Bifunctor (second)
import           Data.Functor ((<&>))
import qualified Control.Exception         as Ex
import qualified Data.HashMap.Strict       as M
import qualified Data.HashSet              as HS
import qualified Data.Char                 as Char
import qualified Data.List                 as L
import qualified Text.PrettyPrint.HughesPJ as PJ

import qualified Language.Fixpoint.Types               as F
-- import qualified Language.Fixpoint.Types.Visitor       as F
import qualified Language.Fixpoint.Misc                as Misc
import           Language.Fixpoint.Types (Expr, ExprV(..), SourcePos) -- , Symbol, symbol)
import qualified Language.Haskell.Liquid.GHC.Misc      as GM
import qualified Liquid.GHC.API       as Ghc
import           Language.Haskell.Liquid.Types.Errors
import           Language.Haskell.Liquid.Types.DataDecl
import           Language.Haskell.Liquid.Types.Names
import qualified Language.Haskell.Liquid.Types.RefType as RT
import           Language.Haskell.Liquid.Types.RType
import           Language.Haskell.Liquid.Types.RTypeOp
import           Language.Haskell.Liquid.Types.Specs
import           Language.Haskell.Liquid.Types.Types
import           Language.Haskell.Liquid.LHNameResolution (symbolToLHName)
import qualified Language.Haskell.Liquid.Misc          as Misc
import qualified Language.Haskell.Liquid.Measure       as Ms
import           Language.Haskell.Liquid.Name.LogicNameEnv (LogicNameEnv(..))
import qualified Language.Haskell.Liquid.Bare.Resolve  as Bare
import qualified Language.Haskell.Liquid.Bare.Types    as Bare
import qualified Language.Haskell.Liquid.Bare.Plugged  as Bare
import           Language.Haskell.Liquid.UX.Config
import qualified Text.Printf                           as Printf

--------------------------------------------------------------------------------
-- | `makeRTEnv` initializes the env needed to `expand` refinements and types,
--   that is, the below needs to be called *before* we use `Expand.expand`
--------------------------------------------------------------------------------
makeRTEnv
  :: LogicNameEnv
  -> ModName
  -> Ms.BareSpec
  -> [(ModName, Ms.BareSpec)]
  -> BareRTEnv
--------------------------------------------------------------------------------
makeRTEnv lenv modName mySpec dependencySpecs
          = renameRTArgs $ makeRTAliases tAs $ makeREAliases eAs
  where
    tAs     = concatMap (Ms.aliases . snd) specs
    eAs     = concatMap (getLHNameExprAliases . snd) specs
    specs = (modName, mySpec) : dependencySpecs

    -- | 'Symbol's are temporarily converted to 'LHName's in expression alias
    -- bodies to use the same lookup and expansion procedure for both
    -- kinds of aliases. Implemented as an specialization of
    -- 'toBareSpecLHName' for the expression aliases field.
    getLHNameExprAliases:: Ms.BareSpec -> [RTAlias F.Symbol (ExprV LHName)]
    getLHNameExprAliases = runIdentity . go

    go :: Ms.BareSpec -> Identity [RTAlias F.Symbol (ExprV LHName)]
    go = mapM (emapRTAlias (\e -> emapExprVM (symToLHName . (++ e)))) . ealiases

    symToLHName = symbolToLHName "makeRTEnv" lenv unhandledNames
    unhandledNames = HS.fromList $ map fst $ expSigs mySpec

-- | We apply @renameRTArgs@ *after* expanding each alias-definition, to
--   ensure that the substitutions work properly (i.e. don't miss expressions
--   hidden inside @RExprArg@ or as strange type parameters.
renameRTArgs :: BareRTEnv -> BareRTEnv
renameRTArgs rte = RTE
  { typeAliases = M.map (renameTys . renameVV . renameRTVArgs) (typeAliases rte)
  , exprAliases = M.map renameRTVArgs (exprAliases rte)
  }

-- | Recursively expands expression aliases by unfolding the definitions of all
--   inner aliases and adds them to the environment.
--   Innermost aliases are unfolded and added first, and an error is thrown if
--   cyclic dependencies are detected.
makeREAliases :: [RTAlias F.Symbol (F.ExprV LHName)] -> BareRTEnv
makeREAliases = graphExpand buildExprEdges f mempty
  where
    f rtEnv xt = setREAlias rtEnv (expand rtEnv (F.loc . rtName $ xt) (lhNametoSymbol xt))
    -- Expression aliases 'LHName's are transformed back to 'Symbol's for the
    -- actual expansion to take place and to be stored in the environment.
    lhNametoSymbol :: RTAlias F.Symbol (F.ExprV LHName) -> RTAlias F.Symbol Expr
    lhNametoSymbol xt = (fmap $ fmap lhNameToResolvedSymbol) xt


-- | @renameTys@ ensures that @RTAlias@ type parameters have distinct names
--   to avoid variable capture e.g. as in T1556.hs
renameTys :: RTAlias F.Symbol BareType -> RTAlias F.Symbol BareType
renameTys rt = rt { rtTArgs = ys, rtBody = sbts (rtBody rt) (zip xs ys) }
  where
    xs    = rtTArgs rt
    ys    = (`F.suffixSymbol` (lhNameToUnqualifiedSymbol . val . rtName $ rt)) <$> xs
    sbts  = foldl (flip subt)


renameVV :: RTAlias F.Symbol BareType -> RTAlias F.Symbol BareType
renameVV rt = rt { rtBody = RT.shiftVV (rtBody rt) (F.vv (Just 0)) }

-- | @renameRTVArgs@ ensures that @RTAlias@ value parameters have distinct names
--   to avoid variable capture e.g. as in tests-names-pos-Capture01.hs
renameRTVArgs :: (F.PPrint a, F.Subable a) => RTAlias x a -> RTAlias x a
renameRTVArgs rt = rt { rtVArgs = newArgs
                      , rtBody  = F.notracepp msg $ F.subst su (rtBody rt)
                      }
  where
    msg          = "renameRTVArgs: " ++ F.showpp su
    su           = F.mkSubst (zip oldArgs (F.eVar <$> newArgs))
    newArgs      = zipWith rtArg (rtVArgs rt) [(0::Int)..]
    oldArgs      = rtVArgs rt
    rtArg x i    = F.suffixSymbol x (F.intSymbol "rta" i)

-- | Recursively expands type aliases by unfolding the definitions of all inner
--   aliases and adds them to the environment.
--   Innermost aliases are unfolded and added first, and an error is thrown if
--   cyclic dependencies are detected.
--   Note that when called from 'makeRTEnv', the input environment contains only
--   expanded expression aliases.
makeRTAliases :: [RTAlias F.Symbol BareType] -> BareRTEnv -> BareRTEnv
makeRTAliases lxts rte = graphExpand buildTypeEdges f rte lxts
  where
    f rtEnv xt = setRTAlias rtEnv (expand rtEnv (F.loc . rtName $ xt) xt)

--------------------------------------------------------------------------------------------------------------

-- | Builds a directed graph of aliases, checks for cyclic dependencies,
--   reorders them so that inner aliases are processed first, and folds over
--   the graph to add each expanded node to the environment.
graphExpand :: (PPrint t)
            => (AliasTable x t -> t -> [LHName])         -- ^ dependencies
            -> (thing -> RTAlias x t -> thing) -- ^ update
            -> thing                                     -- ^ initial
            -> [RTAlias x t]                   -- ^ vertices
            -> thing                                     -- ^ final
graphExpand buildEdges expBody env lxts
           = L.foldl' expBody env (genExpandOrder table' graph)
  where
    table  = buildAliasTable lxts
    graph  = buildAliasGraph (buildEdges table) lxts
    table' = checkCyclicAliases table graph

-- | Inserts a type alias into the environment.
setRTAlias :: RTEnv x t -> RTAlias x t -> RTEnv x t
setRTAlias env a = env { typeAliases =  M.insert n a (typeAliases env) }
  where
    n            = val . rtName $ a

-- | Inserts an expression alias into the environment.
setREAlias :: RTEnv x t -> RTAlias F.Symbol F.Expr -> RTEnv x t
setREAlias env a = env { exprAliases = M.insert n a (exprAliases env) }
  where
    n            = val . rtName $ a

--------------------------------------------------------------------------------

type AliasTable x t = M.HashMap LHName (RTAlias x t)

buildAliasTable :: [RTAlias x t] -> AliasTable x t
buildAliasTable = M.fromList . map (\rta -> (val . rtName $ rta, rta))

fromAliasLHName :: AliasTable x t -> LHName -> RTAlias x t
fromAliasLHName table lhname
  = fromMaybe err (M.lookup lhname table)
  where
    err = panic Nothing ("fromAliasLHName: Dangling alias name: " ++ show lhname)

-- | An adjacency list of nodes representing a directed graph.
--   Used to detect cyclic alias dependencies and to order the expansion
--   of aliases.
type Graph t = [Node t]
-- | A node described by a label, a key, and a list of connected nodes,
--   all parameterized by the same type. This type is used to represent
--   aliases nested within other aliases.
type Node  t = (t, t, [t])

buildAliasGraph :: (PPrint t) => (t -> [LHName]) -> [RTAlias x t]
                -> Graph LHName
buildAliasGraph buildEdges = map (buildNode buildEdges)
  where
    buildNode :: (PPrint t) => (t -> [LHName]) -> RTAlias x t
               -> Node LHName
    buildNode f a = (val . rtName $ a, val . rtName $ a, f (rtBody a))

checkCyclicAliases :: AliasTable x t -> Graph LHName -> AliasTable x t
checkCyclicAliases table graph
  = case mapMaybe go (stronglyConnComp graph) of
      []   -> table
      sccs -> Ex.throw (cycleAliasErr table <$> sccs)
    where
      go (CyclicSCC vs) = Just vs
      go (AcyclicSCC _) = Nothing

cycleAliasErr :: AliasTable x t -> [LHName] -> Error
cycleAliasErr _ []          = panic Nothing "checkCyclicAliases: No aliases in reported cycle"
cycleAliasErr t nameList@(name:_) = ErrAliasCycle { pos    = fst (locate name)
                                                , acycle = map locate nameList }
  where
    locate n = ( GM.fSrcSpan . rtName $ fromAliasLHName t n
                 , pprint n )

-- | Orders aliases so that nested ones are processed first.
genExpandOrder :: AliasTable x t -> Graph LHName -> [RTAlias x t]
genExpandOrder table graph
  = map (fromAliasLHName table) nameOrder
  where
    (digraph, lookupVertex, _)
      = graphFromEdges graph
    nameOrder
      = map (Misc.fst3 . lookupVertex) $ reverse $ topSort digraph

--------------------------------------------------------------------------------

-- | Gathers all constructor names within a the body of a type alias
--   that match a key from the type 'AliasTable'.
buildTypeEdges :: AliasTable x t -> BareType -> [LHName]
buildTypeEdges table = Misc.ordNub . go
  where
    go (RApp c ts rs _) = go_alias (val $ btc_tc c) ++ concatMap go ts ++ concatMap go (mapMaybe go_ref rs)
    go (RFun _ _ t1 t2 _) = go t1 ++ go t2
    go (RAppTy t1 t2 _) = go t1 ++ go t2
    go (RAllE _ t1 t2)  = go t1 ++ go t2
    go (REx _ t1 t2)    = go t1 ++ go t2
    go (RAllT _ t _)    = go t
    go (RAllP _ t)      = go t
    go (RVar _ _)       = []
    go (RExprArg _)     = []
    go (RHole _)        = []
    go (RRTy env _ _ t) = concatMap (go . snd) env ++ go t
    go_alias c          = [c | M.member c table]
    go_ref (RProp _ (RHole _)) = Nothing
    go_ref (RProp  _ t) = Just t

-- | Gathers all variable names within the body of an expression alias
--   that match a key from the expression 'AliasTable'.
buildExprEdges :: AliasTable x t -> F.ExprV LHName -> [LHName]
buildExprEdges table  = Misc.ordNub . go
  where
    go (EApp e1 e2)   = go e1 ++ go e2
    go (ENeg e)       = go e
    go (EBin _ e1 e2) = go e1 ++ go e2
    go (EIte _ e1 e2) = go e1 ++ go e2
    go (ECst e _)     = go e
    go (ESym _)       = []
    go (ECon _)       = []
    go (EVar v)       = go_alias v
    go (PAnd ps)       = concatMap go ps
    go (POr ps)        = concatMap go ps
    go (PNot p)        = go p
    go (PImp p q)      = go p ++ go q
    go (PIff p q)      = go p ++ go q
    go (PAll _ p)      = go p
    go (ELam _ e)      = go e
    go (ELet _ e1 e2)  = go e1 ++ go e2
    go (ECoerc _ _ e)  = go e
    go (PAtom _ e1 e2) = go e1 ++ go e2
    go (ETApp e _)     = go e
    go (ETAbs e _)     = go e
    go (PKVar _ _)     = []
    go (PExist _ e)    = go e
    go_alias f         = [f | M.member f table ]


----------------------------------------------------------------------------------
-- | Using the `BareRTEnv` to do alias-expansion
----------------------------------------------------------------------------------
class Expand a where
  expand :: BareRTEnv -> F.SourcePos -> a -> a

----------------------------------------------------------------------------------
expandLoc :: (Expand a) => BareRTEnv -> Located a -> Located a
expandLoc rtEnv lx = expand rtEnv (F.loc lx) <$> lx

instance Expand Expr where
  expand = expandExpr

instance Expand F.Reft where
  expand rtEnv l (F.Reft (v, ra)) = F.Reft (v, expand rtEnv l ra)

instance Expand RReft where
  expand rtEnv l = fmap (expand rtEnv l)

expandReft :: (Expand r) => BareRTEnv -> F.SourcePos -> RType c tv r -> RType c tv r
expandReft rtEnv l = fmap (expand rtEnv l)
-- expandReft rtEnv l = emapReft (expand rtEnv l)


-- | @expand@ on a SpecType simply expands the refinements,
--   i.e. *does not* apply the type aliases, but just the
--   1. predicate aliases,
--   2. inlines,
--   3. stuff from @LogicMap@

instance Expand SpecType where
  expand = expandReft

-- | @expand@ on a BareType actually applies the type- and expression- aliases.
instance Expand BareType where
  expand rtEnv l
    = expandReft     rtEnv l -- apply expression aliases
    . expandBareType rtEnv l -- apply type       aliases

instance Expand () where
  expand _ _ = id

instance Expand (BRType ()) where
  expand rtEnv l
    = expandReft     rtEnv l -- apply expression aliases
    . void
    . expandBareType rtEnv l -- apply type       aliases
    . fmap (const mempty)

instance Expand (RTAlias F.Symbol Expr) where
  expand rtEnv l x = x { rtBody = expand rtEnv l (rtBody x) }

instance Expand BareRTAlias where
  expand rtEnv l x = x { rtBody = expand rtEnv l (rtBody x) }

instance Expand Body where
  expand rtEnv l (P   p) = P   (expand rtEnv l p)
  expand rtEnv l (E   e) = E   (expand rtEnv l e)
  expand rtEnv l (R x p) = R x (expand rtEnv l p)

instance Expand DataCtor where
  expand rtEnv l c = c
    { dcTheta  = expand rtEnv l (dcTheta c)
    , dcFields = [(x, expand rtEnv l t) | (x, t) <- dcFields c ]
    , dcResult = expand rtEnv l (dcResult c)
    }

instance Expand DataDecl where
  expand rtEnv l d = d
    { tycDCons  = expand rtEnv l (tycDCons  d)
    , tycPropTy = expand rtEnv l (tycPropTy d)
    }

instance Expand BareMeasure where
  expand rtEnv l m = m
    { msSort = expand rtEnv l (msSort m)
    , msEqns = expand rtEnv l (msEqns m)
    }

instance Expand BareDef where
  expand rtEnv l d = d
    { dsort = expand rtEnv l (dsort d)
    , binds = [ (x, expand rtEnv l t) | (x, t) <- binds d]
    , body  = expand rtEnv l (body  d)
    }

instance Expand Ms.BareSpec where
  expand = expandBareSpec

instance Expand a => Expand (F.Located a) where
  expand rtEnv _ = expandLoc rtEnv

instance Expand a => Expand (F.LocSymbol, a) where
  expand rtEnv l (x, y) = (x, expand rtEnv l y)

instance Expand a => Expand (Maybe a) where
  expand rtEnv l = fmap (expand rtEnv l)

instance Expand a => Expand [a] where
  expand rtEnv l = fmap (expand rtEnv l)

instance Expand a => Expand (M.HashMap k a) where
  expand rtEnv l = fmap (expand rtEnv l)

-- | Expands a 'BareSpec'.
expandBareSpec :: BareRTEnv -> F.SourcePos -> Ms.BareSpec -> Ms.BareSpec
expandBareSpec rtEnv l sp = sp
  { measures   = expand rtEnv l (measures   sp)
  , asmSigs    = map (second (expand rtEnv l)) (asmSigs sp)
  , sigs       = map (second (expand rtEnv l)) (sigs sp)
  , ialiases   = [ (f x, f y) | (x, y) <- ialiases sp ]
  , dataDecls  = expand rtEnv l (dataDecls  sp)
  , newtyDecls = expand rtEnv l (newtyDecls sp)
  }
  where f      = expand rtEnv l

expandBareType :: BareRTEnv -> F.SourcePos -> BareType -> BareType
expandBareType rtEnv l = go
  where
    go (RApp c ts rs r)  = case lookupRTEnv c rtEnv of
                             Just rta -> expandRTAliasApp (GM.fSourcePos c) rta (go <$> ts) r
                             Nothing  -> RApp c (go <$> ts) (goRef <$> rs) r
    go (RAppTy t1 t2 r)  = RAppTy (go t1) (go t2) r
    go (RFun  x i t1 t2 r) = RFun  x i (go t1) (go t2) r
    go (RAllT a t r)     = RAllT a (go t) r
    go (RAllP a t)       = RAllP a (go t)
    go (RAllE x t1 t2)   = RAllE x (go t1) (go t2)
    go (REx x t1 t2)     = REx   x (go t1) (go t2)
    go (RRTy e r o t)    = RRTy  e r o     (go t)
    go t@RHole{}         = t
    go t@RVar{}          = t
    go t@RExprArg{}      = t
    goRef (RProp ss t)   = RProp (map (expand rtEnv l <$>) ss) (go t)

lookupRTEnv :: BTyCon -> BareRTEnv -> Maybe BareRTAlias
lookupRTEnv c rtEnv = M.lookup (val $ btc_tc c) (typeAliases rtEnv)

expandRTAliasApp :: F.SourcePos -> BareRTAlias -> [BareType] -> RReft -> BareType
expandRTAliasApp l rta@(RTA {rtName = Loc la _ _}) args r = case isOK of
  Just e     -> Ex.throw e
  Nothing    -> F.subst esu . (`RT.strengthen` r) . RT.subsTyVarsMeet tsu $ rtBody rta
  where
    tsu       = zipWith (\α t -> (α, toRSort t, t)) αs ts
    esu       = F.mkSubst $ zip (F.symbol <$> εs) es
    es        = exprArgFromBareType l msg <$> es0
    (ts, es0) = splitAt nαs args
    (αs, εs)  = (BTV . dummyLoc <$> rtTArgs rta, rtVArgs rta)
    targs     = takeWhile (not . isRExprArg) args
    eargs     = dropWhile (not . isRExprArg) args

    -- ERROR Checking Code
    msg       = "EXPAND-RTALIAS-APP: " ++ F.showpp (rtName rta)
    nαs       = length αs
    nεs       = length εs
    nargs     = length args
    ntargs    = length targs
    neargs    = length eargs
    err       = errRTAliasApp l la rta
    isOK :: Maybe Error
    isOK
      | nargs /= ntargs + neargs
      = err $ PJ.hsep ["Expects", pprint nαs, "type arguments and then", pprint nεs, "expression arguments, but is given", pprint nargs]
      | nargs /= nαs + nεs
      = err $ PJ.hsep ["Expects", pprint nαs, "type arguments and"     , pprint nεs, "expression arguments, but is given", pprint nargs]
      | nαs /= ntargs, not (null eargs)
      = err $ PJ.hsep ["Expects", pprint nαs, "type arguments before expression arguments"]
      | otherwise
      = Nothing

-- | A copy of 'LHNameResolution.exprArg' tailored to the types needed in this
-- module.
exprArgFromBareType :: SourcePos -> String -> BareType -> Expr
exprArgFromBareType l msg = go
  where
    go :: BareType -> Expr
    go (RExprArg e)     = val e
    go (RVar x _)       = EVar $ F.symbol x
    go (RApp x [] [] _) = EVar (getLHNameSymbol $ val $ btc_tc x)
    go (RApp f ts [] _) = F.eApps (EVar (getLHNameSymbol $ val $ btc_tc f)) (go <$> ts)
    go (RAppTy t1 t2 _) = EApp (go t1) (go t2)
    go z                = panic sp $ Printf.printf "Unexpected expression parameter: %s in %s" (show z) msg
    sp                  = Just (GM.sourcePosSrcSpan l)

isRExprArg :: RType c tv r -> Bool
isRExprArg (RExprArg _) = True
isRExprArg _            = False

errRTAliasApp :: F.SourcePos -> F.SourcePos -> BareRTAlias -> PJ.Doc -> Maybe Error
errRTAliasApp l la rta = Just . ErrAliasApp  sp name sp'
  where
    name            = pprint              (rtName rta)
    sp              = GM.sourcePosSrcSpan l
    sp'             = GM.sourcePosSrcSpan la


----------------------------------------------------------------------------------------
-- | @cookSpecType@ is the central place where a @BareType@ gets processed,
--   in multiple steps, into a @SpecType@. See [NOTE:Cooking-SpecType] for
--   details of each of the individual steps.
----------------------------------------------------------------------------------------
cookSpecType :: Bare.Env -> Bare.SigEnv -> ModName -> Bare.PlugTV Ghc.Var -> LocBareType
             -> LocSpecType
cookSpecType env sigEnv name x bt =
  either Ex.throw id $
  cookSpecTypeE env sigEnv name x bt
  where
    _msg = "cookSpecType: " ++ GM.showPpr (z, Ghc.varType <$> z)
    z    = Bare.plugSrc x


-----------------------------------------------------------------------------------------
cookSpecTypeE :: Bare.Env -> Bare.SigEnv -> ModName -> Bare.PlugTV Ghc.Var -> LocBareType
              -> Bare.Lookup LocSpecType
-----------------------------------------------------------------------------------------
cookSpecTypeE env sigEnv name@(ModName _ _) x bt
  = fmap f . bareSpecType env $ bareExpandType rtEnv bt
  where
    f :: LocSpecType -> LocSpecType
    f =   (if doplug || not allowTC then plugHoles allowTC sigEnv name x else id)
        . fmap (RT.addTyConInfo embs tyi)
        . Bare.txRefSort tyi embs
        . fmap txExpToBind -- What does this function DO
        . (specExpandType rtEnv . fmap (generalizeWith x))
        . (if doplug || not allowTC then maybePlug allowTC sigEnv name x else id)

    allowTC = typeclass (getConfig env)
    -- modT   = mname `S.member` wiredInMods
    doplug
      | Bare.LqTV v <- x
      , GM.isMethod v || GM.isSCSel v
      , not (isTarget name)
      = False
      | otherwise
      = True
    _msg i = "cook-" ++ show i ++ " : " ++ F.showpp x
    rtEnv  = Bare.sigRTEnv    sigEnv
    embs   = Bare.sigEmbs     sigEnv
    tyi    = Bare.sigTyRTyMap sigEnv

-- | We don't want to generalize type variables that maybe bound in the
--   outer scope, e.g. see tests/basic/pos/LocalPlug00.hs

generalizeWith :: Bare.PlugTV Ghc.Var -> SpecType -> SpecType
generalizeWith (Bare.HsTV v) t = generalizeVar v t
generalizeWith  Bare.RawTV   t = t
generalizeWith _             t = RT.generalize t

generalizeVar :: Ghc.Var -> SpecType -> SpecType
generalizeVar v t = mkUnivs [(a, mempty) | a <- as] [] t
  where
    as            = filter isGen (RT.freeTyVars t)
    (vas,_)       = Ghc.splitForAllTyCoVars (GM.expandVarType v)
    isGen (RTVar (RTV a) _) = a `elem` vas

-- splitForAllTyCoVars :: Type -> ([TyVar], Type)
--
-- generalize :: (Eq tv) => RType c tv r -> RType c tv r
-- generalize t = mkUnivs (freeTyVars t) [] [] t


bareExpandType :: BareRTEnv -> LocBareType -> LocBareType
bareExpandType = expandLoc

specExpandType :: BareRTEnv -> LocSpecType -> LocSpecType
specExpandType = expandLoc

bareSpecType :: Bare.Env -> LocBareType -> Bare.Lookup LocSpecType
bareSpecType env bt =
  case Bare.ofBareTypeE env (F.loc bt) Nothing (val bt) of
    Left e  -> Left e
    Right t -> Right (F.atLoc bt t)

maybePlug :: Bool -> Bare.SigEnv -> ModName -> Bare.PlugTV Ghc.Var -> LocSpecType -> LocSpecType
maybePlug allowTC sigEnv name kx =
  case Bare.plugSrc kx of
    Nothing -> id
    Just _  -> plugHoles allowTC sigEnv name kx

plugHoles :: Bool -> Bare.SigEnv -> ModName -> Bare.PlugTV Ghc.Var -> LocSpecType -> LocSpecType
plugHoles allowTC sigEnv name =
  Bare.makePluggedSig allowTC name embs tyi exports
  where
    embs    = Bare.sigEmbs     sigEnv
    tyi     = Bare.sigTyRTyMap sigEnv
    exports = Bare.sigExports  sigEnv

{- [NOTE:Cooking-SpecType]
    A @SpecType@ is _raw_ when it is obtained directly from a @BareType@, i.e.
    just by replacing all the @BTyCon@ with @RTyCon@. Before it can be used
    for constraint generation, we need to _cook_ it via the following transforms:

    A @SigEnv@ should contain _all_ the information needed to do the below steps.

    - expand               : resolving all type/refinement etc. aliases
    - ofType               : convert BareType -> SpecType
    - plugged              : filling in any remaining "holes"
    - txRefSort            : filling in the abstract-refinement predicates etc. (YUCK)
    - resolve              : renaming / qualifying symbols?
    - expand (again)       : as the "resolve" step can rename variables to trigger more aliases (e.g. member -> Data.Set.Internal.Member -> Set_mem)
    - generalize           : (universally) quantify free type variables
    - strengthen-measures  : ?
    - strengthen-inline(?) : ?

-}

-----------------------------------------------------------------------------------------------
-- | From BareOLD.Expand
-----------------------------------------------------------------------------------------------


{- TODO-REBARE
instance Expand ty => Expand (Def ty ctor) where
  expand z (Def f xts c t bxts b) =
    Def f <$> expand z xts
          <*> pure c
          <*> expand z t
          <*> expand z bxts
          <*> expand z b

instance Expand ty => Expand (Measure ty ctor) where
  expand z (M n t ds k) =
    M n <$> expand z t <*> expand z ds <*> pure k

instance Expand DataConP where
  expand z d = do
    tyRes'    <- expand z (tyRes     d)
    tyConsts' <- expand z (tyConstrs d)
    tyArgs'   <- expand z (tyArgs    d)
    return d { tyRes =  tyRes', tyConstrs = tyConsts', tyArgs = tyArgs' }
-}

--------------------------------------------------------------------------------
-- | @expandExpr@ applies the aliases and inlines in @BareRTEnv@ to its argument
--   @Expr@. It must first @resolve@ the symbols in the refinement to see if
--   they correspond to alias definitions. However, we ensure that we do not
--   resolve bound variables (e.g. those bound in output refinements by input
--   parameters), and we use the @bs@ parameter to pass in the bound symbols.
--------------------------------------------------------------------------------
expandExpr :: BareRTEnv -> F.SourcePos -> Expr -> Expr
expandExpr rtEnv l      = go
  where
    go e@(EApp _ _)     = expandEApp rtEnv l (F.splitEApp e)
    go (EVar x)         = expandSym  rtEnv l x
    go (ENeg e)         = ENeg       (go e)
    go (ECst e s)       = ECst       (go e) s
    go (PAnd ps)        = PAnd       (go <$> ps)
    go (POr ps)         = POr        (go <$> ps)
    go (PNot p)         = PNot       (go p)
    go (PAll xs p)      = PAll xs    (go p)
    go (PExist xs p)    = PExist xs  (go p)
    go (ELam xt e)      = ELam xt    (go e)
    go (ELet x e1 e2)   = ELet x     (go e1) (go e2)
    go (ECoerc a t e)   = ECoerc a t (go e)
    go (ETApp e s)      = ETApp      (go e) s
    go (ETAbs e s)      = ETAbs      (go e) s
    go (EBin op e1 e2)  = EBin op    (go e1) (go e2)
    go (PImp    e1 e2)  = PImp       (go e1) (go e2)
    go (PIff    e1 e2)  = PIff       (go e1) (go e2)
    go (PAtom b e1 e2)  = PAtom b    (go e1) (go e2)
    go (EIte  p e1 e2)  = EIte (go p)(go e1) (go e2)
    go e@(PKVar _ _)    = e
    go e@(ESym _)       = e
    go e@(ECon _)       = e

expandSym :: BareRTEnv -> F.SourcePos -> F.Symbol -> Expr
expandSym rtEnv l s' = expandEApp rtEnv l (EVar s', [])

-- REBARE :: expandSym' :: Symbol -> BareM Symbol
-- REBARE :: expandSym' s = do
  -- REBARE :: axs <- gets axSyms
  -- REBARE :: let s' = dropModuleNamesAndUnique s
  -- REBARE :: return $ if M.member s' axs then s' else s

expandEApp :: BareRTEnv -> F.SourcePos -> (Expr, [Expr]) -> Expr
expandEApp rtEnv l (EVar f, es) = case mBody of
    Just re -> expandApp l   re       es'
    Nothing -> F.eApps       (EVar f) es'
  where
    eAs     = M.mapKeys lhNameToResolvedSymbol $ exprAliases rtEnv
    mBody   = M.lookup f eAs `mplus` M.lookup (GM.dropModuleUnique f) eAs
    es'     = expandExpr rtEnv l <$> es
    _f0     = GM.dropModuleNamesAndUnique f

expandEApp _ _ (f, es) = F.eApps f es

--------------------------------------------------------------------------------
-- | Expand Alias Application --------------------------------------------------
--------------------------------------------------------------------------------
expandApp :: F.Subable ty => F.SourcePos -> RTAlias F.Symbol ty -> [Expr] -> ty
expandApp l re es
  | Just su <- args = F.subst su (rtBody re)
  | otherwise       = Ex.throw err
  where
    args            = F.mkSubst <$> Misc.zipMaybe (rtVArgs re) es
    err             :: UserError
    err             = ErrAliasApp sp alias sp' msg
    sp              = GM.sourcePosSrcSpan l
    alias           = pprint           (rtName re)
    sp'             = GM.fSrcSpan . rtName $ re
    msg             =  "expects" PJ.<+> pprint (length $ rtVArgs re)
                   PJ.<+> "arguments but it is given"
                   PJ.<+> pprint (length es)


-------------------------------------------------------------------------------
-- | Replace Predicate Arguments With Existentials ----------------------------
-------------------------------------------------------------------------------
txExpToBind   :: SpecType -> SpecType
-------------------------------------------------------------------------------
txExpToBind t =
  evalState (expToBindT t) (ExSt 0 M.empty πs)
  where
    πs        = M.fromList [(pname p, p) | p <- ty_preds $ toRTypeRep t ]

data ExSt = ExSt { fresh :: Int
                 , emap  :: M.HashMap F.Symbol (RSort, F.Expr)
                 , pmap  :: M.HashMap F.Symbol RPVar
                 }

-- | TODO: Niki please write more documentation for this, maybe an example?
--   I can't really tell whats going on... (RJ)

expToBindT :: SpecType -> State ExSt SpecType
expToBindT (RVar v r)
  = expToBindRef r >>= addExists . RVar v
expToBindT (RFun x i t1 t2 r)
  = do t1' <- expToBindT t1
       t2' <- expToBindT t2
       expToBindRef r >>= addExists . RFun x i t1' t2'
expToBindT (RAllT a t r)
  = do t' <- expToBindT t
       expToBindRef r >>= addExists . RAllT a t'
expToBindT (RAllP p t)
  = fmap (RAllP p) (expToBindT t)
expToBindT (RApp c ts rs r)
  = do ts' <- mapM expToBindT ts
       rs' <- mapM expToBindReft rs
       expToBindRef r >>= addExists . RApp c ts' rs'
expToBindT (RAppTy t1 t2 r)
  = do t1' <- expToBindT t1
       t2' <- expToBindT t2
       expToBindRef r >>= addExists . RAppTy t1' t2'
expToBindT (RRTy xts r o t)
  = do xts' <- zip xs <$> mapM expToBindT ts
       r'   <- expToBindRef r
       t'   <- expToBindT t
       return $ RRTy xts' r' o t'
  where
     (xs, ts) = unzip xts
expToBindT t
  = return t

expToBindReft              :: SpecProp -> State ExSt SpecProp
expToBindReft (RProp s (RHole r)) = rPropP s <$> expToBindRef r
expToBindReft (RProp s t)  = RProp s  <$> expToBindT t


getBinds :: State ExSt (M.HashMap F.Symbol (RSort, F.Expr))
getBinds
  = do bds <- gets emap
       modify $ \st -> st{emap = M.empty}
       return bds

addExists :: SpecType -> State ExSt SpecType
addExists t = fmap (M.foldlWithKey' addExist t) getBinds

addExist :: SpecType -> F.Symbol -> (RSort, F.Expr) -> SpecType
addExist t x (tx, e) = REx x t' t
  where
    t'               = ofRSort tx `RT.strengthen` RT.uTop r
    r                = F.exprReft e

expToBindRef :: UReft r -> State ExSt (UReft r)
expToBindRef (MkUReft r (Pr p))
  = mapM expToBind p <&> (MkUReft r . Pr)

expToBind :: UsedPVar -> State ExSt UsedPVar
expToBind p = do
  res <- gets (M.lookup (pname p) . pmap)
  case res of
    Nothing ->
      panic Nothing ("expToBind: " ++ show p)
    Just π  -> do
      let pargs0 = zip (pargs p) (Misc.fst3 <$> pargs π)
      pargs' <- mapM expToBindParg pargs0
      return $ p { pargs = pargs' }

expToBindParg :: (((), F.Symbol, F.Expr), RSort) -> State ExSt ((), F.Symbol, F.Expr)
expToBindParg ((t, s, e), s') = fmap ((,,) t s) (expToBindExpr e s')

expToBindExpr :: F.Expr ->  RSort -> State ExSt F.Expr
expToBindExpr e@(EVar s) _
  | Char.isLower $ F.headSym $ F.symbol s
  = return e
expToBindExpr e t
  = do s <- freshSymbol
       modify $ \st -> st{emap = M.insert s (t, e) (emap st)}
       return $ EVar s

freshSymbol :: State ExSt F.Symbol
freshSymbol
  = do n <- gets fresh
       modify $ \s -> s {fresh = n+1}
       return $ F.symbol $ "ex#" ++ show n


-- wiredInMods :: S.HashSet Ghc.ModuleName
-- wiredInMods = S.fromList $ Ghc.mkModuleName <$>
--   ["Language.Haskell.Liquid.String",
--   "Language.Haskell.Liquid.Prelude",
--   "Language.Haskell.Liquid.Foreign",
--   "Language.Haskell.Liquid.Bag",
--   "Prelude",
--   "System.IO",
--   "Data.Word",
--   "Data.Time.Calendar",
--   "Data.Set",
--   "Data.Either",
--   "Data.ByteString.Unsafe",
--   "Data.ByteString.Lazy",
--   "Data.ByteString.Short",
--   "Data.Foldable",
--   "Data.OldList",
--   "Data.Text",
--   "Data.Tuple",
--   "Data.Bits",
--   "Data.Chare",
--   "Data.String",
--   "Data.Vector",
--   "Data.Time",
--   "Data.Int",
--   "Data.Text.Fusion",
--   "Data.Map",
--   "Data.Text.Fusion.Common",
--   "KMeansHelper",
--   "Data.Text.Lazy.Fusion",
--   "Control.Exception",
--   "Control.Parallel.Strategies",
--   "Data.Traversable",
--   "GHC.Read",
--   "Data.ByteString",
--   "GHC.Classes",
--   "GHC.Ptr",
--   "GHC.Word",
--   "Language.Haskell.Liquid.Equational",
--   "GHC.Types",
--   "GHC.Num",
--   "GHC.CString",
--   "GHC.IO.Handle",
--   "GHC.Prim",
--   "GHC.Int",
--   "GHC.Base",
--   "Foreign.Ptr",
--   "GHC.ForeignPtr",
--   "GHC.List",
--   "Foreign.C.String",
--   "GHC.Exts",
--   "Foreign.Marshal.Alloc",
--   "Foreign.Marshal.Array",
--   "Foreign.C.Types",
--   "GHC.Real",
--   "Foreign.Storable",
--   "Foreign.ForeignPtr"]