content: {-# LANGUAGE FlexibleContexts      #-}
  {-# LANGUAGE FlexibleInstances     #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE OverloadedStrings     #-}
  {-# LANGUAGE ScopedTypeVariables   #-}
  {-# LANGUAGE TupleSections         #-}
  {-# LANGUAGE UndecidableInstances  #-}
  {-# LANGUAGE ConstraintKinds       #-}
  
  module Language.Haskell.Liquid.Types.Fresh
    ( Freshable(..)
    , refreshTy
    , refreshVV
    , refreshArgs
    , refreshHoles
    , refreshArgsSub
    )
    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           CoreUtils  (exprType)
  import           Prelude                        hiding (error)
  -- import           Type       (Type)
  -- import           CoreSyn
  -- import           Var        (varType, isTyVar, Var)
  
  import qualified Language.Fixpoint.Types as F
  -- import           Language.Fixpoint.Types.Visitor (kvars)
  import           Language.Haskell.Liquid.Misc  (single)
  import           Language.Haskell.Liquid.Types.Errors
  import           Language.Haskell.Liquid.Types.RType
  import           Language.Haskell.Liquid.Types.RTypeOp
  import           Language.Haskell.Liquid.Types.Types
  import           Language.Haskell.Liquid.Types.RefType
  
  
  class (Applicative m, Monad m) => Freshable m a where
    fresh   :: m a
    true    :: Bool -> a -> m a
    true _  = return
    refresh :: Bool -> a -> m a
    refresh _ = return
  
  
  instance (Freshable m Integer, Monad m, Applicative m) => Freshable m F.Symbol where
    fresh = F.tempSymbol "x" <$> fresh
  
  instance (Freshable m Integer, Monad m, Applicative m) => Freshable m F.Expr where
    fresh  = kv <$> fresh
      where
        kv = (`F.PKVar` mempty) . F.intKvar
  
  instance (Freshable m Integer, Monad m, Applicative m) => Freshable m [F.Expr] where
    fresh = single <$> fresh
  
  instance (Freshable m Integer, Monad m, Applicative m) => Freshable m F.Reft where
    fresh                  = panic Nothing "fresh Reft"
    true    _ (F.Reft (v,_)) = return $ F.Reft (v, F.PTrue)
    refresh _ (F.Reft (_,_)) = (F.Reft .) . (,) <$> freshVV <*> fresh
      where
        freshVV            = F.vv . Just <$> fresh
  
  instance Freshable m Integer => Freshable m RReft where
    fresh             = panic Nothing "fresh RReft"
    true allowTC (MkUReft r _)    = MkUReft <$> true allowTC r    <*> return mempty
    refresh allowTC (MkUReft r _) = MkUReft <$> refresh allowTC r <*> return mempty
  
  instance (Freshable m Integer, Freshable m r, Reftable r ) => Freshable m (RRType r) where
    fresh   = panic Nothing "fresh RefType"
    refresh = refreshRefType
    true    = trueRefType
  
  -----------------------------------------------------------------------------------------------
  trueRefType :: (Freshable m Integer, Freshable m r, Reftable r) => Bool -> RRType r -> m (RRType r)
  -----------------------------------------------------------------------------------------------
  trueRefType allowTC (RAllT α t r)
    = RAllT α <$> true allowTC t <*> true allowTC r
  
  trueRefType allowTC (RAllP π t)
    = RAllP π <$> true allowTC t
  
  trueRefType allowTC (RFun _ _ t t' _)
    -- YL: attaching rfinfo here is crucial
    = rFun' (classRFInfo allowTC) <$> fresh <*> true allowTC t <*> true allowTC t'
  
  trueRefType allowTC (RApp c ts _  _) | if allowTC then isEmbeddedDict c else isClass c
    = rRCls c <$> mapM (true allowTC) ts
  
  trueRefType allowTC (RApp c ts rs r)
    = RApp c <$> mapM (true allowTC) ts <*> mapM (trueRef allowTC) rs <*> true allowTC r
  
  trueRefType allowTC (RAppTy t t' _)
    = RAppTy <$> true allowTC t <*> true allowTC t' <*> return mempty
  
  trueRefType allowTC (RVar a r)
    = RVar a <$> true allowTC r
  
  trueRefType allowTC (RAllE y ty tx)
    = do y'  <- fresh
         ty' <- true allowTC ty
         tx' <- true allowTC tx
         return $ RAllE y' ty' (tx' `F.subst1` (y, F.EVar y'))
  
  trueRefType allowTC (RRTy e o r t)
    = RRTy e o r <$> trueRefType allowTC t
  
  trueRefType allowTC (REx _ t t')
    = REx <$> fresh <*> true allowTC t <*> true allowTC t'
  
  trueRefType _ t@(RExprArg _)
    = return t
  
  trueRefType _ t@(RHole _)
    = return t
  
  trueRef :: (Reftable r, Freshable f r, Freshable f Integer)
          => Bool -> Ref τ (RType RTyCon RTyVar r) -> f (Ref τ (RRType r))
  trueRef _ (RProp _ (RHole _)) = panic Nothing "trueRef: unexpected RProp _ (RHole _))"
  trueRef allowTC (RProp s t) = RProp s <$> trueRefType allowTC t
  
  
  -----------------------------------------------------------------------------------------------
  refreshRefType :: (Freshable m Integer, Freshable m r, Reftable r) => Bool -> RRType r -> m (RRType r)
  -----------------------------------------------------------------------------------------------
  refreshRefType allowTC (RAllT α t r)
    = RAllT α <$> refresh allowTC t <*> true allowTC r
  
  refreshRefType allowTC (RAllP π t)
    = RAllP π <$> refresh allowTC t
  
  refreshRefType allowTC (RFun sym i t t' _)
    | sym == F.dummySymbol = (\b t1 t2 -> RFun b i t1 t2 mempty) <$> fresh <*> refresh allowTC t <*> refresh allowTC t'
    | otherwise          = (\t1 t2 -> RFun sym i t1 t2 mempty)   <$> refresh allowTC t <*> refresh allowTC t'
  
  refreshRefType _ (RApp rc ts _ _) | isClass rc
    = return $ rRCls rc ts
  
  refreshRefType allowTC (RApp rc ts rs r)
    = RApp rc <$> mapM (refresh allowTC) ts <*> mapM (refreshRef allowTC) rs <*> refresh allowTC r
  
  refreshRefType allowTC (RVar a r)
    = RVar a <$> refresh allowTC r
  
  refreshRefType allowTC (RAppTy t t' r)
    = RAppTy <$> refresh allowTC t <*> refresh allowTC t' <*> refresh allowTC r
  
  refreshRefType allowTC (RAllE y ty tx)
    = do y'  <- fresh
         ty' <- refresh allowTC ty
         tx' <- refresh allowTC tx
         return $ RAllE y' ty' (tx' `F.subst1` (y, F.EVar y'))
  
  refreshRefType allowTC (RRTy e o r t)
    = RRTy e o r <$> refreshRefType allowTC t
  
  refreshRefType _ t
    = return t
  
  refreshRef :: (Reftable r, Freshable f r, Freshable f Integer)
             => Bool -> Ref τ (RType RTyCon RTyVar r) -> f (Ref τ (RRType r))
  refreshRef _ (RProp _ (RHole _)) = panic Nothing "refreshRef: unexpected (RProp _ (RHole _))"
  refreshRef allowTC (RProp s t) = RProp <$> mapM freshSym s <*> refreshRefType allowTC t
  
  freshSym :: Freshable f a => (t, t1) -> f (a, t1)
  freshSym (_, t)        = (, t) <$> fresh
  
  
  --------------------------------------------------------------------------------
  refreshTy :: (FreshM m) => SpecType -> m SpecType
  --------------------------------------------------------------------------------
  refreshTy t = refreshVV t >>= refreshArgs
  
  --------------------------------------------------------------------------------
  type FreshM m = Freshable m Integer
  --------------------------------------------------------------------------------
  
  --------------------------------------------------------------------------------
  refreshVV :: FreshM m => SpecType -> m SpecType
  --------------------------------------------------------------------------------
  refreshVV (RAllT a t r) =
    RAllT a <$> refreshVV t <*> return r
  
  refreshVV (RAllP p t) =
    RAllP p <$> refreshVV t
  
  refreshVV (REx x t1 t2) = do
    t1' <- refreshVV t1
    t2' <- refreshVV t2
    shiftVV (REx x t1' t2') <$> fresh
  
  refreshVV (RFun x i t1 t2 r) = do
    t1' <- refreshVV t1
    t2' <- refreshVV t2
    shiftVV (RFun x i t1' t2' r) <$> fresh
  
  refreshVV (RAppTy t1 t2 r) = do
    t1' <- refreshVV t1
    t2' <- refreshVV t2
    shiftVV (RAppTy t1' t2' r) <$> fresh
  
  refreshVV (RApp c ts rs r) = do
    ts' <- mapM refreshVV    ts
    rs' <- mapM refreshVVRef rs
    shiftVV (RApp c ts' rs' r) <$> fresh
  
  refreshVV t =
    shiftVV t <$> fresh
  
  refreshVVRef :: Freshable m Integer => Ref b SpecType -> m (Ref b SpecType)
  refreshVVRef (RProp ss (RHole r))
    = return $ RProp ss (RHole r)
  
  refreshVVRef (RProp ss t)
    = do xs    <- mapM (const fresh) syms
         let su = F.mkSubst $ zip syms (F.EVar <$> xs)
         t'    <- refreshVV t
         return $ RProp (zip xs bs) (F.subst su t')
      where
      (syms, bs) = unzip ss
  
  --------------------------------------------------------------------------------
  refreshArgs :: (FreshM m) => SpecType -> m SpecType
  --------------------------------------------------------------------------------
  refreshArgs t = fst <$> refreshArgsSub t
  
  
  -- NV TODO: this does not refresh args if they are wrapped in an RRTy
  refreshArgsSub :: (FreshM m) => SpecType -> m (SpecType, F.Subst)
  refreshArgsSub t
    = do ts     <- mapM refreshArgs ts_u
         xs'    <- mapM (const fresh) xs
         let sus = F.mkSubst <$> L.inits (zip xs (F.EVar <$> xs'))
         let su  = last sus
         ts'    <- mapM refreshPs $ zipWith F.subst sus ts
         let rs' = zipWith F.subst sus rs
         tr     <- refreshPs $ F.subst su tbd
         let t'  = fromRTypeRep $ trep {ty_binds = xs', ty_args = ts', ty_res = tr, ty_refts = rs'}
         return (t', su)
      where
         trep    = toRTypeRep t
         xs      = ty_binds trep
         ts_u    = ty_args  trep
         tbd     = ty_res   trep
         rs      = ty_refts trep
  
  refreshPs :: (FreshM m) => SpecType -> m SpecType
  refreshPs = mapPropM go
    where
      go (RProp s st) = do
        t'    <- refreshPs st
        xs    <- mapM (const fresh) s
        let su = F.mkSubst [(y, F.EVar x) | (x, (y, _)) <- zip xs s]
        return $ RProp [(x, t) | (x, (_, t)) <- zip xs s] $ F.subst su t'
  
  --------------------------------------------------------------------------------
  refreshHoles :: (F.Symbolic t, Reftable r, TyConable c, Freshable f r)
               => Bool -> [(t, RType c tv r)] -> f ([F.Symbol], [(t, RType c tv r)])
  refreshHoles allowTC vts = first catMaybes . unzip . map extract <$> mapM (refreshHoles' allowTC) vts
    where
    --   extract :: (t, t1, t2) -> (t, (t1, t2))
      extract (a,b,c) = (a,(b,c))
  
  refreshHoles' :: (F.Symbolic a, Reftable r, TyConable c, Freshable m r)
                => Bool -> (a, RType c tv r) -> m (Maybe F.Symbol, a, RType c tv r)
  refreshHoles' allowTC (x,t)
    | noHoles t = return (Nothing, x, t)
    | otherwise = (Just $ F.symbol x,x,) <$> mapReftM tx t
    where
      tx r | hasHole r = refresh allowTC r
           | otherwise = return r
  
  noHoles :: (Reftable r, TyConable c) => RType c tv r -> Bool
  noHoles = and . foldReft False (\_ r bs -> not (hasHole r) : bs) []
  