-- Copyright 2022 Google LLC
--
-- Use of this source code is governed by a BSD-style
-- license that can be found in the LICENSE file or at
-- https://developers.google.com/open-source/licenses/bsd

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}

module Subst where

import Control.Applicative
import Control.Monad.Identity
import Control.Monad.Reader
import Control.Monad.State.Strict

import Name
import MTL1
import IRVariants
import Types.Core
import Types.Top
import Core
import qualified RawName as R
import QueryTypePure
import Err

-- === SubstReader class ===

class (SinkableV v, Monad2 m) => SubstReader (v::V) (m::MonadKind2) | m -> v where
   getSubst :: m i o (Subst v i o)
   withSubst :: Subst v i' o -> m i' o a -> m i o a

lookupSubstM :: (Color c, SubstReader v m) => Name c i -> m i o (v c o)
lookupSubstM name = (!name) <$> getSubst

dropSubst :: (SubstReader v m, FromName v) => m o o a -> m i o a
dropSubst cont = withSubst idSubst cont
{-# INLINE dropSubst #-}

withVoidSubst :: (SubstReader v m, FromName v) => m VoidS o a -> m i o a
withVoidSubst cont = withSubst (newSubst absurdNameFunction) cont
{-# INLINE withVoidSubst #-}

extendSubst :: SubstReader v m => SubstFrag v i i' o -> m i' o a -> m i o a
extendSubst frag cont = do
  env <- (<>>frag) <$> getSubst
  withSubst env cont
{-# INLINE extendSubst #-}

-- XXX: this only (monadically) visits each name once, even if a name has
-- multiple occurrences. So don't use it to count occurrences or anything like
-- that! It's not deliberate. It's just an accident of the implementation, where
-- we gather the (de-duplicated) free names and then traverse them. At some
-- point we may add a monadic traversal to `Subst{E,B}`, which would actually
-- visit each occurrence.
traverseNames
  :: forall v e m i o.
     (SubstE v e, HoistableE e, SinkableE e, FromName v, EnvReader m)
  => (forall c. Color c => Name c i -> m o (v c o))
  -> e i -> m o (e o)
traverseNames f e = do
  let vs = freeVarsE e
  m <- flip R.traverseWithKey (fromNameSet vs) \rawName (SubstItem d c _) ->
    interpretColor c \(ColorProxy :: ColorProxy c) -> do
      v' <- f (UnsafeMakeName rawName :: Name c i)
      return $ SubstItem d c (unsafeCoerceVC v')
  fmapNamesM (applyTraversed m) e
{-# INLINE traverseNames #-}

applyTraversed :: FromName v
               => RawNameMap (SubstItem v n) -> Name c i -> v c n
applyTraversed m = \((UnsafeMakeName v) :: Name c i) -> case R.lookup v m of
    Just item -> unsafeFromSubstItem item
    Nothing -> fromName $ (UnsafeMakeName v :: Name c o)

fmapNames :: (SubstE v e, Distinct o)
          => Env o -> (forall c. Color c => Name c i -> v c o) -> e i -> e o
fmapNames env f e = substE (env, newSubst f) e
{-# INLINE fmapNames #-}

fmapNamesM :: (SubstE v e, SinkableE e, EnvReader m)
          => (forall c. Color c => Name c i -> v c o)
          -> e i -> m o (e o)
fmapNamesM f e = do
  env <- unsafeGetEnv
  Distinct <- getDistinct
  return $ substE (env, newSubst f) e
{-# INLINE fmapNamesM #-}

-- === type classes for traversing names ===

class FromName v => SubstE (v::V) (e::E) where
  -- TODO: can't make an alias for these constraints because of impredicativity
  substE :: Distinct o => (Env o, Subst v i o) -> e i -> e o

  default substE :: (GenericE e, SubstE v (RepE e), Distinct o)
                 => (Env o, Subst v i o) -> e i -> e o
  substE env e = toE $ substE env (fromE e)

class (FromName v, SinkableB b) => SubstB (v::V) (b::B) where
  substB
    :: Distinct o
    => (Env o, Subst v i o)
    -> b i i'
    -> (forall o'. Distinct o' => (Env o', Subst v i' o') -> b o o' -> a)
    -> a

  default substB
    :: (GenericB b, SubstB v (RepB b))
    => Distinct o
    => (Env o, Subst v i o)
    -> b i i'
    -> (forall o'. Distinct o' => (Env o', Subst v i' o') -> b o o' -> a)
    -> a
  substB env b cont =
    substB env (fromB b) \env' b' ->
      cont env' $ toB b'

class ( FromName substVal, SinkableV v
      , forall c. Color c => SubstE substVal (v c))
      => SubstV (substVal::V) (v::V)

instance ( forall r. IRRep r => SinkableE (atom r)
         , forall r. IRRep r => RenameE (atom r)) => RenameV (SubstVal atom)

instance (Color c, forall r. IRRep r => RenameE (atom r)) => RenameE (SubstVal atom c) where
  renameE (_, env) (Rename name) = Rename $ env ! name
  renameE (scope, env) (SubstVal atom) = SubstVal $ renameE (scope, env) atom

substM :: (SubstReader v m, EnvReader2 m, SinkableE e, SubstE v e, FromName v)
       => e i -> m i o (e o)
substM e = do
  subst <- getSubst
  substM' subst e
{-# INLINE substM #-}

substM' :: (EnvReader m, SinkableE e, SubstE v e, FromName v)
       => Subst v i o -> e i -> m o (e o)
substM' subst e = do
  case tryApplyIdentitySubst subst e of
    Just e' -> return $ e'
    Nothing -> do
      env <- unsafeGetEnv
      Distinct <- getDistinct
      return $ fmapNames env (subst!) e
{-# INLINE substM' #-}

fromConstAbs :: (BindsNames b, HoistableE e) => Abs b e n -> HoistExcept (e n)
fromConstAbs (Abs b e) = hoist b e

-- === rename-only substitutions ===

extendRenamer :: (SubstReader v m, FromName v) => SubstFrag Name i i' o -> m i' o r -> m i o r
extendRenamer frag = extendSubst (fmapSubstFrag (const fromName) frag)

extendBinderRename
  :: (SubstReader v m, FromName v, BindsAtMostOneName b c, BindsOneName b' c)
  => b i i' -> b' o o' -> m i' o' r -> m i o' r
extendBinderRename b b' cont = extendSubst (b@>fromName (binderName b')) cont
{-# INLINE extendBinderRename #-}

applyRename
  :: (ScopeReader m, RenameE e, SinkableE e)
  => Ext h o => SubstFrag Name h i o -> e i -> m o (e o)
applyRename substFrag x = do
  Distinct <- getDistinct
  scope <- unsafeGetScope
  let subst = sink idSubst <>> substFrag
  case tryApplyIdentitySubst subst x of
    Just x' -> return x'
    Nothing -> return $ renameE (scope, newSubst (subst!)) x
{-# INLINE applyRename #-}

renameM
  :: (SubstReader Name m, ScopeReader2 m, SinkableE e, RenameE e)
  => e i -> m i o (e o)
renameM e = do
  env <- getSubst
  case tryApplyIdentitySubst env e of
    Just e' -> return $ e'
    Nothing -> do
      WithScope scope env' <- addScope env
      sinkM $ renameE (scope, newSubst (env'!)) e
{-# INLINE renameM #-}

renameBinders
  :: (EnvExtender2 m, SubstReader Name m, RenameB b, BindsEnv b)
  => b i i'
  -> (forall o'. DExt o o' => b o o' -> m i' o' a)
  -> m i o a
renameBinders b cont = do
  ab <- renameM $ Abs b $ idSubstFrag b
  refreshAbs ab \b' subst -> extendSubst subst $ cont b'
{-# INLINE renameBinders #-}

-- === various convenience utilities ===

applySubstFragPure :: (SubstE v e, SinkableE e, SinkableV v, FromName v, Ext h o, Distinct o)
                   => Env o -> SubstFrag v h i o -> e i -> e o
applySubstFragPure env substFrag x = do
  let fullSubst = sink idSubst <>> substFrag
  applySubstPure env fullSubst x

applySubstPure :: (SubstE v e, SinkableE e, SinkableV v, FromName v, Distinct o)
               => Env o -> Subst v i o -> e i -> e o
applySubstPure env subst x = do
  case tryApplyIdentitySubst subst x of
    Just x' -> x'
    Nothing -> fmapNames env (subst !) x

applySubst :: (EnvReader m, SubstE v e, SinkableE e, SinkableV v, FromName v)
           => Ext h o
           => SubstFrag v h i o -> e i -> m o (e o)
applySubst substFrag x = do
  Distinct <- getDistinct
  env <- unsafeGetEnv
  return $ applySubstFragPure env substFrag x
{-# INLINE applySubst #-}

applyAbs :: ( SinkableV v, SinkableE e
            , FromName v, EnvReader m, BindsOneName b c, SubstE v e)
         => Abs b e n -> v c n -> m n (e n)
applyAbs (Abs b body) x = applySubst (b@>x) body
{-# INLINE applyAbs #-}

applyNaryAbs :: ( SinkableV v, FromName v, EnvReader m, BindsNameList b c, SubstE v e
                , SubstB v b, SinkableE e)
             => Abs b e n -> [v c n] -> m n (e n)
applyNaryAbs (Abs bs body) xs = applySubst (bs @@> xs) body
{-# INLINE applyNaryAbs #-}

substBinders
  :: ( SinkableV v, SubstV v v, EnvExtender2 m, FromName v
     , SubstReader v m, SubstB v b, RenameV v, RenameB b, BindsEnv b)
  => b i i'
  -> (forall o'. DExt o o' => b o o' -> m i' o' a)
  -> m i o a
substBinders b cont = do
  substBindersFrag b \subst b' -> extendSubst subst $ cont b'
{-# INLINE substBinders #-}

substBindersFrag
  :: ( SinkableV v, SubstV v v, EnvExtender2 m, FromName v
     , SubstReader v m, SubstB v b, RenameV v, RenameB b, BindsEnv b)
  => b i i'
  -> (forall o'. DExt o o' => SubstFrag v i i' o' -> b o o' -> m i o' a)
  -> m i o a
substBindersFrag b cont = do
  ab <- substM $ Abs b $ idSubstFrag b
  refreshAbs ab \b' subst -> cont subst b'
{-# INLINE substBindersFrag #-}

-- === atom subst vals ===

data SubstVal (atom::IR->E) (c::C) (n::S) where
  SubstVal :: IRRep r => atom r n -> SubstVal atom (AtomNameC r) n
  Rename   :: Name c n -> SubstVal atom c n
type AtomSubstVal = SubstVal Atom

type family IsAtomName (c::C) where
  IsAtomName (AtomNameC r) = True
  IsAtomName _             = False

instance (Color c, IsAtomName c ~ False) => SubstE (SubstVal atom) (Name c) where
  substE (_, env) v = case env ! v of Rename v' -> v'

instance FromName (SubstVal atom) where
  fromName = Rename
  {-# INLINE fromName #-}

class ToSubstVal (v::V) (atom::IR->E) where
  toSubstVal :: v c n -> SubstVal atom c n

instance ToSubstVal (Name::V) (atom::IR->E) where
  toSubstVal = Rename

instance ToSubstVal (SubstVal atom) atom where
  toSubstVal = id

type AtomSubstReader v m = (SubstReader v m, FromName v, ToSubstVal v Atom)

toAtomVar :: (EnvReader m,  IRRep r) => AtomName r n -> m n (AtomVar r n)
toAtomVar v = do
  ty <- getType <$> lookupAtomName v
  return $ AtomVar v ty

lookupAtomSubst :: (IRRep r, SubstReader AtomSubstVal m, EnvReader2 m) => AtomName r i -> m i o (Atom r o)
lookupAtomSubst v = do
  lookupSubstM v >>= \case
    Rename v' -> toAtom <$> toAtomVar v'
    SubstVal x -> return x

atomSubstM :: (AtomSubstReader v m, EnvReader2 m, SinkableE e, SubstE AtomSubstVal e)
           => e i -> m i o (e o)
atomSubstM e = do
  subst <- getSubst
  let subst' = asAtomSubstValSubst subst
  substM' subst' e

asAtomSubstValSubst :: ToSubstVal v Atom => Subst v i o -> Subst AtomSubstVal i o
asAtomSubstValSubst subst = newSubst \v -> toSubstVal (subst ! v)

-- === SubstReaderT transformer ===

newtype SubstReaderT (v::V) (m::MonadKind1) (i::S) (o::S) (a:: *) =
  SubstReaderT' { runSubstReaderT' :: ReaderT (Subst v i o) (m o) a }

pattern SubstReaderT :: (Subst v i o -> m o a) -> SubstReaderT v m i o a
pattern SubstReaderT f = SubstReaderT' (ReaderT f)

runSubstReaderT :: Subst v i o -> SubstReaderT v m i o a -> m o a
runSubstReaderT env m = runReaderT (runSubstReaderT' m) env
{-# INLINE runSubstReaderT #-}

instance (forall n. Functor (m n)) => Functor (SubstReaderT v m i o) where
  fmap f (SubstReaderT' m) = SubstReaderT' $ fmap f m
  {-# INLINE fmap #-}

instance Monad1 m => Applicative (SubstReaderT v m i o) where
  pure   = SubstReaderT' . pure
  {-# INLINE pure #-}
  liftA2 f (SubstReaderT' x) (SubstReaderT' y) = SubstReaderT' $ liftA2 f x y
  {-# INLINE liftA2 #-}
  (SubstReaderT' f) <*> (SubstReaderT' x) = SubstReaderT' $ f <*> x
  {-# INLINE (<*>) #-}

instance (forall n. Monad (m n)) => Monad (SubstReaderT v m i o) where
  return = SubstReaderT' . return
  {-# INLINE return #-}
  (SubstReaderT' m) >>= f = SubstReaderT' (m >>= (runSubstReaderT' . f))
  {-# INLINE (>>=) #-}

deriving instance (Monad1 m, MonadFail1   m) => MonadFail   (SubstReaderT v m i o)
deriving instance (Monad1 m, Alternative1 m) => Alternative (SubstReaderT v m i o)
deriving instance Fallible1  m => Fallible  (SubstReaderT v m i o)
deriving instance Catchable1 m => Catchable (SubstReaderT v m i o)

type ScopedSubstReader (v::V) = SubstReaderT v (ScopeReaderT Identity) :: MonadKind2

liftSubstReaderT :: Monad1 m => m o a -> SubstReaderT v m i o a
liftSubstReaderT m = SubstReaderT' $ lift m
{-# INLINE liftSubstReaderT #-}

runScopedSubstReader :: Distinct o => Scope o -> Subst v i o
                   -> ScopedSubstReader v i o a -> a
runScopedSubstReader scope env m =
  runIdentity $ runScopeReaderT scope $ runSubstReaderT env m
{-# INLINE runScopedSubstReader #-}

withSubstReaderT :: FromName v => SubstReaderT v m n n a -> m n a
withSubstReaderT = runSubstReaderT idSubst
{-# INLINE withSubstReaderT #-}

instance (SinkableV v, Monad1 m) => SubstReader v (SubstReaderT v m) where
  getSubst = SubstReaderT' ask
  {-# INLINE getSubst #-}
  withSubst env (SubstReaderT' cont) = SubstReaderT' $ withReaderT (const env) cont
  {-# INLINE withSubst #-}

instance (SinkableV v, ScopeReader m) => ScopeReader (SubstReaderT v m i) where
  unsafeGetScope = liftSubstReaderT unsafeGetScope
  {-# INLINE unsafeGetScope #-}
  getDistinct = liftSubstReaderT getDistinct
  {-# INLINE getDistinct #-}

instance (SinkableV v, EnvReader m) => EnvReader (SubstReaderT v m i) where
  unsafeGetEnv = liftSubstReaderT unsafeGetEnv
  {-# INLINE unsafeGetEnv #-}

instance (SinkableV v, ScopeReader m, EnvExtender m)
         => EnvExtender (SubstReaderT v m i) where
  refreshAbs ab cont = SubstReaderT \subst ->
    refreshAbs ab \b e -> do
      subst' <- sinkM subst
      let SubstReaderT cont' = cont b e
      cont' subst'
  {-# INLINE refreshAbs #-}

instance MonadDiffState1 m s d => MonadDiffState1 (SubstReaderT v m i) s d where
  withDiffState s m =
    SubstReaderT \subst -> do
      withDiffState s $ runSubstReaderT subst m

  updateDiffStateM d = liftSubstReaderT $ updateDiffStateM d
  getDiffState = liftSubstReaderT getDiffState

type SubstEnvReaderM v = SubstReaderT v EnvReaderM :: MonadKind2

liftSubstEnvReaderM
  :: forall v m n a. (EnvReader m, FromName v)
  => SubstEnvReaderM v n n a
  -> m n a
liftSubstEnvReaderM cont = liftEnvReaderM $ runSubstReaderT idSubst $ cont
{-# INLINE liftSubstEnvReaderM #-}

instance (SinkableV v, ScopeReader m, ScopeExtender m)
         => ScopeExtender (SubstReaderT v m i) where
  refreshAbsScope ab cont = SubstReaderT \env ->
    refreshAbsScope ab \b e -> do
      let SubstReaderT cont' = cont b e
      env' <- sinkM env
      cont' env'

instance (SinkableV v, MonadIO1 m) => MonadIO (SubstReaderT v m i o) where
  liftIO m = liftSubstReaderT $ liftIO m
  {-# INLINE liftIO #-}

instance (Monad1 m, MonadState (s o) (m o)) => MonadState (s o) (SubstReaderT v m i o) where
  state = liftSubstReaderT . state
  {-# INLINE state #-}

instance (Monad1 m, MonadReader (r o) (m o)) => MonadReader (r o) (SubstReaderT v m i o) where
  ask = SubstReaderT $ const ask
  {-# INLINE ask #-}
  local r (SubstReaderT' (ReaderT f)) = SubstReaderT \env -> local r $ f env
  {-# INLINE local #-}

-- === instances ===

instance (forall r. IRRep r => SinkableE (atom r)) => SinkableV (SubstVal atom)
instance (forall r. IRRep r => SinkableE (atom r)) => SinkableE (SubstVal atom c) where
  sinkingProofE fresh substVal = case substVal of
    Rename name  -> Rename   $ sinkingProofE fresh name
    SubstVal val -> SubstVal $ sinkingProofE fresh val

instance (SubstB v b, SubstE v e) => SubstE v (Abs b e) where
  substE env (Abs b body) = do
    substB env b \env' b' -> Abs b' $ substE env' body

instance ( BindsNames b1, SubstB v b1
         , BindsNames b2, SubstB v b2
         , SinkableV v, FromName v)
         => SubstB v (PairB b1 b2) where
  substB env (PairB b1 b2) cont =
    substB env b1 \env' b1' ->
      substB env' b2 \env'' b2' ->
        cont env'' $ PairB b1' b2'

instance (SinkableE e, SubstE v e) => SubstB v (LiftB e) where
  substB env@(_, subst) (LiftB e) cont = case tryApplyIdentitySubst subst e of
    Just e' -> cont env $ LiftB e'
    Nothing -> cont env $ LiftB $ substE env e
  {-# INLINE substB #-}

instance (SubstB v b1, SubstB v b2) => SubstB v (EitherB b1 b2) where
  substB env (LeftB b) cont =
    substB env b \env' b' ->
      cont env' $ LeftB b'
  substB env (RightB b) cont =
    substB env b \env' b' ->
      cont env' $ RightB b'

instance (Color c, SinkableE ann, SubstE v ann, SinkableV v, ToBinding ann c)
         => SubstB v (BinderP c ann) where
  substB (env, subst) (b:>ty) cont = do
    let ty' = substE (env, subst) ty
    withFresh (getNameHint b) (toScope env) \bRaw -> do
      let b' = bRaw:>ty'
      let env' = env `extendOutMap` toEnvFrag b'
      let UnsafeMakeName bn  = binderName b
      let UnsafeMakeName bn' = binderName b'
      let subst' = case subst of
                   UnsafeMakeIdentitySubst | bn == bn' -> UnsafeMakeIdentitySubst
                   _ -> sink subst <>> b @> (fromName $ binderName b')
      cont (env', subst') b'

instance (BindsNames b, SubstB v b, SinkableV v)
         => SubstB v (Nest b) where
  substB env (Nest b bs) cont =
    substB env b \env' b' ->
      substB env' bs \env'' bs' ->
        cont env'' $ Nest b' bs'
  substB env Empty cont = cont env Empty

instance FromName v => SubstE v UnitE where
  substE _ UnitE = UnitE

instance SubstB v b => SubstB v (WithAttrB a b) where
  substB env (WithAttrB x b) cont =
    substB env b \env' b' -> cont env' $ WithAttrB x b'

instance (Traversable f, SubstE v e) => SubstE v (ComposeE f e) where
  substE env (ComposeE xs) = ComposeE $ fmap (substE env) xs

instance (SubstE v e1, SubstE v e2) => SubstE v (PairE e1 e2) where
  substE env (PairE x y) = PairE (substE env x) (substE env y)

instance (SubstE v e1, SubstE v e2) => SubstE v (EitherE e1 e2) where
  substE env (LeftE  x) = LeftE  $ substE env x
  substE env (RightE x) = RightE $ substE env x

instance FromName v => SubstE v VoidE where
  substE _ _ = error "impossible"

instance FromName v => SubstE v (LiftE a) where
  substE _ (LiftE x) = LiftE x

instance SubstE v e => SubstE v (ListE e) where
  substE env (ListE xs) = ListE $ map (substE env) xs

instance SubstE v e => SubstE v (RListE e) where
  substE env (RListE xs) = RListE $ fmap (substE env) xs

instance SubstE v e => SubstE v (NonEmptyListE e) where
  substE env (NonEmptyListE xs) = NonEmptyListE $ fmap (substE env) xs

instance (p ~ True => SubstE v e, FromName v) => SubstE v (WhenE p e) where
  substE (scope, subst) (WhenE e) = WhenE $ substE (scope, subst) e

instance (r ~ r' => SubstE v e, FromName v) => SubstE v (WhenIRE r r' e) where
  substE (scope, subst) (WhenIRE e) = WhenIRE $ substE (scope, subst) e

instance SubstV substVal v => SubstE substVal (SubstFrag v i i') where
   substE env frag = fmapSubstFrag (\_ val -> substE env val) frag

instance SubstV substVal v => SubstE substVal (Subst v i) where
  substE env = \case
    Subst f frag -> Subst (\n -> substE env (f n)) $ substE env frag
    UnsafeMakeIdentitySubst
      -> Subst (\n -> substE env (fromName $ unsafeCoerceE n)) emptyInFrag

instance (SubstE v e0, SubstE v e1, SubstE v e2,
          SubstE v e3, SubstE v e4, SubstE v e5,
          SubstE v e6, SubstE v e7)
            => SubstE v (EitherE8 e0 e1 e2 e3 e4 e5 e6 e7) where
  substE env = \case
    Case0 e -> Case0 $ substE env e
    Case1 e -> Case1 $ substE env e
    Case2 e -> Case2 $ substE env e
    Case3 e -> Case3 $ substE env e
    Case4 e -> Case4 $ substE env e
    Case5 e -> Case5 $ substE env e
    Case6 e -> Case6 $ substE env e
    Case7 e -> Case7 $ substE env e
  {-# INLINE substE #-}