-- editorconfig-checker-disable-file
{-# LANGUAGE LambdaCase          #-}
{-# LANGUAGE ScopedTypeVariables #-}
module UntypedPlutusCore.Subst
    ( substVarA
    , substVar
    , termSubstNamesM
    , termSubstNames
    , termMapNames
    , programMapNames
    , vTerm
    ) where

import PlutusPrelude

import UntypedPlutusCore.Core

import Control.Lens

purely :: ((a -> Identity b) -> c -> Identity d) -> (a -> b) -> c -> d
purely :: forall a b c d.
((a -> Identity b) -> c -> Identity d) -> (a -> b) -> c -> d
purely = ((a -> Identity b) -> c -> Identity d) -> (a -> b) -> c -> d
forall a b. Coercible a b => a -> b
coerce

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

-- | Replace a variable using the given function.
substVar
    :: (name -> Maybe (Term name uni fun ann))
    -> Term name uni fun ann
    -> Term name uni fun ann
substVar :: forall name (uni :: * -> *) fun ann.
(name -> Maybe (Term name uni fun ann))
-> Term name uni fun ann -> Term name uni fun ann
substVar = ((name -> Identity (Maybe (Term name uni fun ann)))
 -> Term name uni fun ann -> Identity (Term name uni fun ann))
-> (name -> Maybe (Term name uni fun ann))
-> Term name uni fun ann
-> Term name uni fun ann
forall a b c d.
((a -> Identity b) -> c -> Identity d) -> (a -> b) -> c -> d
purely (name -> Identity (Maybe (Term name uni fun ann)))
-> Term name uni fun ann -> Identity (Term name uni fun ann)
forall (f :: * -> *) name (uni :: * -> *) fun ann.
Applicative f =>
(name -> f (Maybe (Term name uni fun ann)))
-> Term name uni fun ann -> f (Term name uni fun ann)
substVarA

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

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

-- | Completely replace the names with a new name type.
termMapNames
    :: forall name name' uni fun ann
    . (name -> name')
    -> Term name uni fun ann
    -> Term name' uni fun ann
termMapNames :: forall name name' (uni :: * -> *) fun ann.
(name -> name') -> Term name uni fun ann -> Term name' uni fun ann
termMapNames name -> name'
f = Term name uni fun ann -> Term name' uni fun ann
go
    where
        -- This is all a bit clunky because of the type-changing, I'm not sure of a nicer way to do it
        go :: Term name uni fun ann -> Term name' uni fun ann
        go :: Term name uni fun ann -> Term name' uni fun ann
go = \case
            LamAbs ann
ann name
name Term name uni fun ann
body -> ann -> name' -> Term name' uni fun ann -> Term name' uni fun ann
forall name (uni :: * -> *) fun ann.
ann -> name -> Term name uni fun ann -> Term name uni fun ann
LamAbs ann
ann (name -> name'
f name
name) (Term name uni fun ann -> Term name' uni fun ann
go Term name uni fun ann
body)
            Var ann
ann name
name         -> ann -> name' -> Term name' uni fun ann
forall name (uni :: * -> *) fun ann.
ann -> name -> Term name uni fun ann
Var ann
ann (name -> name'
f name
name)

            Apply ann
ann Term name uni fun ann
t1 Term name uni fun ann
t2      -> ann
-> Term name' uni fun ann
-> Term name' uni fun ann
-> Term name' uni fun ann
forall name (uni :: * -> *) fun ann.
ann
-> Term name uni fun ann
-> Term name uni fun ann
-> Term name uni fun ann
Apply ann
ann (Term name uni fun ann -> Term name' uni fun ann
go Term name uni fun ann
t1) (Term name uni fun ann -> Term name' uni fun ann
go Term name uni fun ann
t2)
            Delay ann
ann Term name uni fun ann
t          -> ann -> Term name' uni fun ann -> Term name' uni fun ann
forall name (uni :: * -> *) fun ann.
ann -> Term name uni fun ann -> Term name uni fun ann
Delay ann
ann (Term name uni fun ann -> Term name' uni fun ann
go Term name uni fun ann
t)
            Force ann
ann Term name uni fun ann
t          -> ann -> Term name' uni fun ann -> Term name' uni fun ann
forall name (uni :: * -> *) fun ann.
ann -> Term name uni fun ann -> Term name uni fun ann
Force ann
ann (Term name uni fun ann -> Term name' uni fun ann
go Term name uni fun ann
t)
            Constr ann
ann Word64
i [Term name uni fun ann]
es      -> ann -> Word64 -> [Term name' uni fun ann] -> Term name' uni fun ann
forall name (uni :: * -> *) fun ann.
ann -> Word64 -> [Term name uni fun ann] -> Term name uni fun ann
Constr ann
ann Word64
i ((Term name uni fun ann -> Term name' uni fun ann)
-> [Term name uni fun ann] -> [Term name' uni fun ann]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Term name uni fun ann -> Term name' uni fun ann
go [Term name uni fun ann]
es)
            Case ann
ann Term name uni fun ann
arg Vector (Term name uni fun ann)
cs      -> ann
-> Term name' uni fun ann
-> Vector (Term name' uni fun ann)
-> Term name' uni fun ann
forall name (uni :: * -> *) fun ann.
ann
-> Term name uni fun ann
-> Vector (Term name uni fun ann)
-> Term name uni fun ann
Case ann
ann (Term name uni fun ann -> Term name' uni fun ann
go Term name uni fun ann
arg) ((Term name uni fun ann -> Term name' uni fun ann)
-> Vector (Term name uni fun ann)
-> Vector (Term name' uni fun ann)
forall a b. (a -> b) -> Vector a -> Vector b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Term name uni fun ann -> Term name' uni fun ann
go Vector (Term name uni fun ann)
cs)

            Constant ann
ann Some (ValueOf uni)
c       -> ann -> Some (ValueOf uni) -> Term name' uni fun ann
forall name (uni :: * -> *) fun ann.
ann -> Some (ValueOf uni) -> Term name uni fun ann
Constant ann
ann Some (ValueOf uni)
c
            Builtin ann
ann fun
b        -> ann -> fun -> Term name' uni fun ann
forall name (uni :: * -> *) fun ann.
ann -> fun -> Term name uni fun ann
Builtin ann
ann fun
b
            Error ann
ann            -> ann -> Term name' uni fun ann
forall name (uni :: * -> *) fun ann. ann -> Term name uni fun ann
Error ann
ann

programMapNames
    :: forall name name' uni fun ann
    . (name -> name')
    -> Program name uni fun ann
    -> Program name' uni fun ann
programMapNames :: forall name name' (uni :: * -> *) fun ann.
(name -> name')
-> Program name uni fun ann -> Program name' uni fun ann
programMapNames name -> name'
f (Program ann
a Version
v Term name uni fun ann
term) = ann
-> Version -> Term name' uni fun ann -> Program name' uni fun ann
forall name (uni :: * -> *) fun ann.
ann -> Version -> Term name uni fun ann -> Program name uni fun ann
Program ann
a Version
v ((name -> name') -> Term name uni fun ann -> Term name' uni fun ann
forall name name' (uni :: * -> *) fun ann.
(name -> name') -> Term name uni fun ann -> Term name' uni fun ann
termMapNames name -> name'
f Term name uni fun ann
term)

-- TODO: this could be a Traversal
-- | Get all the term variables in a term.
vTerm :: Fold (Term name uni fun ann) name
vTerm :: forall name (uni :: * -> *) fun ann (f :: * -> *).
(Contravariant f, Applicative f) =>
(name -> f name)
-> Term name uni fun ann -> f (Term name uni fun ann)
vTerm = (Term name uni fun ann -> f (Term name uni fun ann))
-> Term name uni fun ann -> f (Term name uni fun ann)
forall name (uni :: * -> *) fun ann (f :: * -> *).
(Contravariant f, Applicative f) =>
(Term name uni fun ann -> f (Term name uni fun ann))
-> Term name uni fun ann -> f (Term name uni fun ann)
termSubtermsDeep ((Term name uni fun ann -> f (Term name uni fun ann))
 -> Term name uni fun ann -> f (Term name uni fun ann))
-> ((name -> f name)
    -> Term name uni fun ann -> f (Term name uni fun ann))
-> (name -> f name)
-> Term name uni fun ann
-> f (Term name uni fun ann)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (name -> f name)
-> Term name uni fun ann -> f (Term name uni fun ann)
forall name (uni :: * -> *) fun ann (f :: * -> *).
Applicative f =>
(name -> f name)
-> Term name uni fun ann -> f (Term name uni fun ann)
termVars