{-# LANGUAGE FlexibleContexts #-}

-- | Implements naive substitution functions for replacing type and term variables.
module PlutusIR.Transform.Substitute (
  substVarA,
  substTyVarA,
  typeSubstTyNames,
  termSubstNames,
  termSubstNamesM,
  termSubstTyNames,
  termSubstTyNamesM,
  bindingSubstNames,
  bindingSubstTyNames,
) where

import PlutusCore.Subst (purely, typeSubstTyNames)
import PlutusIR
import PlutusPrelude

import Control.Lens

-- | Applicatively replace a variable using the given function.
substVarA ::
  (Applicative f) =>
  (name -> f (Maybe (Term tyname name uni fun ann))) ->
  Term tyname name uni fun ann ->
  f (Term tyname name uni fun ann)
substVarA :: forall (f :: * -> *) name tyname (uni :: * -> *) fun ann.
Applicative f =>
(name -> f (Maybe (Term tyname name uni fun ann)))
-> Term tyname name uni fun ann -> f (Term tyname name uni fun ann)
substVarA name -> f (Maybe (Term tyname name uni fun ann))
nameF t :: Term tyname name uni fun ann
t@(Var ann
_ name
name) = Term tyname name uni fun ann
-> Maybe (Term tyname name uni fun ann)
-> Term tyname name uni fun ann
forall a. a -> Maybe a -> a
fromMaybe Term tyname name uni fun ann
t (Maybe (Term tyname name uni fun ann)
 -> Term tyname name uni fun ann)
-> f (Maybe (Term tyname name uni fun ann))
-> f (Term tyname name uni fun ann)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> name -> f (Maybe (Term tyname name uni fun ann))
nameF name
name
substVarA name -> f (Maybe (Term tyname name uni fun ann))
_ Term tyname name uni fun ann
t                  = Term tyname name uni fun ann -> f (Term tyname name uni fun ann)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Term tyname name uni fun ann
t
{-# INLINE substVarA #-}

-- | Applicatively replace a type variable using the given function.
substTyVarA ::
  (Applicative f) =>
  (tyname -> f (Maybe (Type tyname uni ann))) ->
  Type tyname uni ann ->
  f (Type tyname uni ann)
substTyVarA :: forall (f :: * -> *) tyname (uni :: * -> *) ann.
Applicative f =>
(tyname -> f (Maybe (Type tyname uni ann)))
-> Type tyname uni ann -> f (Type tyname uni ann)
substTyVarA tyname -> f (Maybe (Type tyname uni ann))
tynameF ty :: Type tyname uni ann
ty@(TyVar ann
_ tyname
tyname) = Type tyname uni ann
-> Maybe (Type tyname uni ann) -> Type tyname uni ann
forall a. a -> Maybe a -> a
fromMaybe Type tyname uni ann
ty (Maybe (Type tyname uni ann) -> Type tyname uni ann)
-> f (Maybe (Type tyname uni ann)) -> f (Type tyname uni ann)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> tyname -> f (Maybe (Type tyname uni ann))
tynameF tyname
tyname
substTyVarA tyname -> f (Maybe (Type tyname uni ann))
_ Type tyname uni ann
ty                        = Type tyname uni ann -> f (Type tyname uni ann)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Type tyname uni ann
ty
{-# INLINE substTyVarA #-}

-- | Naively substitute names using the given functions (i.e. do not substitute binders).
termSubstNames ::
  (name -> Maybe (Term tyname name uni fun a)) ->
  Term tyname name uni fun a ->
  Term tyname name uni fun a
termSubstNames :: forall name tyname (uni :: * -> *) fun a.
(name -> Maybe (Term tyname name uni fun a))
-> Term tyname name uni fun a -> Term tyname name uni fun a
termSubstNames = ((name -> Identity (Maybe (Term tyname name uni fun a)))
 -> Term tyname name uni fun a
 -> Identity (Term tyname name uni fun a))
-> (name -> Maybe (Term tyname name uni fun a))
-> Term tyname name uni fun a
-> Term tyname name uni fun a
forall a b c d.
((a -> Identity b) -> c -> Identity d) -> (a -> b) -> c -> d
purely (name -> Identity (Maybe (Term tyname name uni fun a)))
-> Term tyname name uni fun a
-> Identity (Term tyname name uni fun a)
forall (m :: * -> *) name tyname (uni :: * -> *) fun ann.
Monad m =>
(name -> m (Maybe (Term tyname name uni fun ann)))
-> Term tyname name uni fun ann -> m (Term tyname name uni fun ann)
termSubstNamesM

-- | Naively monadically substitute names using the given function (i.e. do not substitute binders).
termSubstNamesM ::
  (Monad m) =>
  (name -> m (Maybe (Term tyname name uni fun ann))) ->
  Term tyname name uni fun ann ->
  m (Term tyname name uni fun ann)
termSubstNamesM :: forall (m :: * -> *) name tyname (uni :: * -> *) fun ann.
Monad m =>
(name -> m (Maybe (Term tyname name uni fun ann)))
-> Term tyname name uni fun ann -> m (Term tyname name uni fun ann)
termSubstNamesM = LensLike
  (WrappedMonad m)
  (Term tyname name uni fun ann)
  (Term tyname name uni fun ann)
  (Term tyname name uni fun ann)
  (Term tyname name uni fun ann)
-> (Term tyname name uni fun ann
    -> m (Term tyname name uni fun ann))
-> Term tyname name uni fun ann
-> m (Term tyname name uni fun ann)
forall (m :: * -> *) a b.
Monad m =>
LensLike (WrappedMonad m) a b a b -> (b -> m b) -> a -> m b
transformMOf LensLike
  (WrappedMonad m)
  (Term tyname name uni fun ann)
  (Term tyname name uni fun ann)
  (Term tyname name uni fun ann)
  (Term tyname name uni fun ann)
forall tyname name (uni :: * -> *) fun a (f :: * -> *).
Applicative f =>
(Term tyname name uni fun a -> f (Term tyname name uni fun a))
-> Term tyname name uni fun a -> f (Term tyname name uni fun a)
termSubterms ((Term tyname name uni fun ann -> m (Term tyname name uni fun ann))
 -> Term tyname name uni fun ann
 -> m (Term tyname name uni fun ann))
-> ((name -> m (Maybe (Term tyname name uni fun ann)))
    -> Term tyname name uni fun ann
    -> m (Term tyname name uni fun ann))
-> (name -> m (Maybe (Term tyname name uni fun ann)))
-> Term tyname name uni fun ann
-> m (Term tyname name uni fun ann)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (name -> m (Maybe (Term tyname name uni fun ann)))
-> Term tyname name uni fun ann -> m (Term tyname name uni fun ann)
forall (f :: * -> *) name tyname (uni :: * -> *) fun ann.
Applicative f =>
(name -> f (Maybe (Term tyname name uni fun ann)))
-> Term tyname name uni fun ann -> f (Term tyname name uni fun ann)
substVarA

-- | Naively substitute type names using the given functions (i.e. do not substitute binders).
termSubstTyNames ::
  (tyname -> Maybe (Type tyname uni a)) ->
  Term tyname name uni fun a ->
  Term tyname name uni fun a
termSubstTyNames :: forall tyname (uni :: * -> *) a name fun.
(tyname -> Maybe (Type tyname uni a))
-> Term tyname name uni fun a -> Term tyname name uni fun a
termSubstTyNames = ((tyname -> Identity (Maybe (Type tyname uni a)))
 -> Term tyname name uni fun a
 -> Identity (Term tyname name uni fun a))
-> (tyname -> Maybe (Type tyname uni a))
-> Term tyname name uni fun a
-> Term tyname name uni fun a
forall a b c d.
((a -> Identity b) -> c -> Identity d) -> (a -> b) -> c -> d
purely (tyname -> Identity (Maybe (Type tyname uni a)))
-> Term tyname name uni fun a
-> Identity (Term tyname name uni fun a)
forall (m :: * -> *) tyname (uni :: * -> *) ann name fun.
Monad m =>
(tyname -> m (Maybe (Type tyname uni ann)))
-> Term tyname name uni fun ann -> m (Term tyname name uni fun ann)
termSubstTyNamesM

{- | Naively monadically substitute type names using the given function
(i.e. do not substitute binders).
-}
termSubstTyNamesM ::
  (Monad m) =>
  (tyname -> m (Maybe (Type tyname uni ann))) ->
  Term tyname name uni fun ann ->
  m (Term tyname name uni fun ann)
termSubstTyNamesM :: forall (m :: * -> *) tyname (uni :: * -> *) ann name fun.
Monad m =>
(tyname -> m (Maybe (Type tyname uni ann)))
-> Term tyname name uni fun ann -> m (Term tyname name uni fun ann)
termSubstTyNamesM =
  LensLike
  (WrappedMonad m)
  (Term tyname name uni fun ann)
  (Term tyname name uni fun ann)
  (Term tyname name uni fun ann)
  (Term tyname name uni fun ann)
-> (Term tyname name uni fun ann
    -> m (Term tyname name uni fun ann))
-> Term tyname name uni fun ann
-> m (Term tyname name uni fun ann)
forall (m :: * -> *) a b.
Monad m =>
LensLike (WrappedMonad m) a b a b -> (b -> m b) -> a -> m b
transformMOf LensLike
  (WrappedMonad m)
  (Term tyname name uni fun ann)
  (Term tyname name uni fun ann)
  (Term tyname name uni fun ann)
  (Term tyname name uni fun ann)
forall tyname name (uni :: * -> *) fun a (f :: * -> *).
Applicative f =>
(Term tyname name uni fun a -> f (Term tyname name uni fun a))
-> Term tyname name uni fun a -> f (Term tyname name uni fun a)
termSubterms ((Term tyname name uni fun ann -> m (Term tyname name uni fun ann))
 -> Term tyname name uni fun ann
 -> m (Term tyname name uni fun ann))
-> ((tyname -> m (Maybe (Type tyname uni ann)))
    -> Term tyname name uni fun ann
    -> m (Term tyname name uni fun ann))
-> (tyname -> m (Maybe (Type tyname uni ann)))
-> Term tyname name uni fun ann
-> m (Term tyname name uni fun ann)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LensLike
  m
  (Term tyname name uni fun ann)
  (Term tyname name uni fun ann)
  (Type tyname uni ann)
  (Type tyname uni ann)
-> LensLike
     m
     (Term tyname name uni fun ann)
     (Term tyname name uni fun ann)
     (Type tyname uni ann)
     (Type tyname uni ann)
forall (f :: * -> *) s t a b.
LensLike f s t a b -> LensLike f s t a b
traverseOf LensLike
  m
  (Term tyname name uni fun ann)
  (Term tyname name uni fun ann)
  (Type tyname uni ann)
  (Type tyname uni ann)
forall tyname name (uni :: * -> *) fun a (f :: * -> *).
Applicative f =>
(Type tyname uni a -> f (Type tyname uni a))
-> Term tyname name uni fun a -> f (Term tyname name uni fun a)
termSubtypes LensLike
  m
  (Term tyname name uni fun ann)
  (Term tyname name uni fun ann)
  (Type tyname uni ann)
  (Type tyname uni ann)
-> ((tyname -> m (Maybe (Type tyname uni ann)))
    -> Type tyname uni ann -> m (Type tyname uni ann))
-> (tyname -> m (Maybe (Type tyname uni ann)))
-> Term tyname name uni fun ann
-> m (Term tyname name uni fun ann)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. LensLike
  (WrappedMonad m)
  (Type tyname uni ann)
  (Type tyname uni ann)
  (Type tyname uni ann)
  (Type tyname uni ann)
-> (Type tyname uni ann -> m (Type tyname uni ann))
-> Type tyname uni ann
-> m (Type tyname uni ann)
forall (m :: * -> *) a b.
Monad m =>
LensLike (WrappedMonad m) a b a b -> (b -> m b) -> a -> m b
transformMOf LensLike
  (WrappedMonad m)
  (Type tyname uni ann)
  (Type tyname uni ann)
  (Type tyname uni ann)
  (Type tyname uni ann)
forall tyname (uni :: * -> *) ann (f :: * -> *).
Applicative f =>
(Type tyname uni ann -> f (Type tyname uni ann))
-> Type tyname uni ann -> f (Type tyname uni ann)
typeSubtypes ((Type tyname uni ann -> m (Type tyname uni ann))
 -> Type tyname uni ann -> m (Type tyname uni ann))
-> ((tyname -> m (Maybe (Type tyname uni ann)))
    -> Type tyname uni ann -> m (Type tyname uni ann))
-> (tyname -> m (Maybe (Type tyname uni ann)))
-> Type tyname uni ann
-> m (Type tyname uni ann)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (tyname -> m (Maybe (Type tyname uni ann)))
-> Type tyname uni ann -> m (Type tyname uni ann)
forall (f :: * -> *) tyname (uni :: * -> *) ann.
Applicative f =>
(tyname -> f (Maybe (Type tyname uni ann)))
-> Type tyname uni ann -> f (Type tyname uni ann)
substTyVarA

-- | Naively substitute names using the given functions (i.e. do not substitute binders).
bindingSubstNames ::
  (name -> Maybe (Term tyname name uni fun a)) ->
  Binding tyname name uni fun a ->
  Binding tyname name uni fun a
bindingSubstNames :: forall name tyname (uni :: * -> *) fun a.
(name -> Maybe (Term tyname name uni fun a))
-> Binding tyname name uni fun a -> Binding tyname name uni fun a
bindingSubstNames name -> Maybe (Term tyname name uni fun a)
nameF = ASetter
  (Binding tyname name uni fun a)
  (Binding tyname name uni fun a)
  (Term tyname name uni fun a)
  (Term tyname name uni fun a)
-> (Term tyname name uni fun a -> Term tyname name uni fun a)
-> Binding tyname name uni fun a
-> Binding tyname name uni fun a
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over ASetter
  (Binding tyname name uni fun a)
  (Binding tyname name uni fun a)
  (Term tyname name uni fun a)
  (Term tyname name uni fun a)
forall tyname name (uni :: * -> *) fun a (f :: * -> *).
Applicative f =>
(Term tyname name uni fun a -> f (Term tyname name uni fun a))
-> Binding tyname name uni fun a
-> f (Binding tyname name uni fun a)
bindingSubterms ((name -> Maybe (Term tyname name uni fun a))
-> Term tyname name uni fun a -> Term tyname name uni fun a
forall name tyname (uni :: * -> *) fun a.
(name -> Maybe (Term tyname name uni fun a))
-> Term tyname name uni fun a -> Term tyname name uni fun a
termSubstNames name -> Maybe (Term tyname name uni fun a)
nameF)

-- | Naively substitute type names using the given functions (i.e. do not substitute binders).
bindingSubstTyNames ::
  (tyname -> Maybe (Type tyname uni a)) ->
  Binding tyname name uni fun a ->
  Binding tyname name uni fun a
bindingSubstTyNames :: forall tyname (uni :: * -> *) a name fun.
(tyname -> Maybe (Type tyname uni a))
-> Binding tyname name uni fun a -> Binding tyname name uni fun a
bindingSubstTyNames tyname -> Maybe (Type tyname uni a)
tynameF =
  ASetter
  (Binding tyname name uni fun a)
  (Binding tyname name uni fun a)
  (Term tyname name uni fun a)
  (Term tyname name uni fun a)
-> (Term tyname name uni fun a -> Term tyname name uni fun a)
-> Binding tyname name uni fun a
-> Binding tyname name uni fun a
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over ASetter
  (Binding tyname name uni fun a)
  (Binding tyname name uni fun a)
  (Term tyname name uni fun a)
  (Term tyname name uni fun a)
forall tyname name (uni :: * -> *) fun a (f :: * -> *).
Applicative f =>
(Term tyname name uni fun a -> f (Term tyname name uni fun a))
-> Binding tyname name uni fun a
-> f (Binding tyname name uni fun a)
bindingSubterms ((tyname -> Maybe (Type tyname uni a))
-> Term tyname name uni fun a -> Term tyname name uni fun a
forall tyname (uni :: * -> *) a name fun.
(tyname -> Maybe (Type tyname uni a))
-> Term tyname name uni fun a -> Term tyname name uni fun a
termSubstTyNames tyname -> Maybe (Type tyname uni a)
tynameF)
    (Binding tyname name uni fun a -> Binding tyname name uni fun a)
-> (Binding tyname name uni fun a -> Binding tyname name uni fun a)
-> Binding tyname name uni fun a
-> Binding tyname name uni fun a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ASetter
  (Binding tyname name uni fun a)
  (Binding tyname name uni fun a)
  (Type tyname uni a)
  (Type tyname uni a)
-> (Type tyname uni a -> Type tyname uni a)
-> Binding tyname name uni fun a
-> Binding tyname name uni fun a
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
over ASetter
  (Binding tyname name uni fun a)
  (Binding tyname name uni fun a)
  (Type tyname uni a)
  (Type tyname uni a)
forall tyname name (uni :: * -> *) fun a (f :: * -> *).
Applicative f =>
(Type tyname uni a -> f (Type tyname uni a))
-> Binding tyname name uni fun a
-> f (Binding tyname name uni fun a)
bindingSubtypes ((tyname -> Maybe (Type tyname uni a))
-> Type tyname uni a -> Type tyname uni a
forall tyname (uni :: * -> *) ann.
(tyname -> Maybe (Type tyname uni ann))
-> Type tyname uni ann -> Type tyname uni ann
typeSubstTyNames tyname -> Maybe (Type tyname uni a)
tynameF)