{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}

module UntypedPlutusCore.Core.Plated
    ( termBinds
    , termVars
    , termUniques
    , termSubterms
    , termSubtermsDeep
    , termUniquesDeep
    ) where

import PlutusCore.Core (HasUniques)
import PlutusCore.Name
import UntypedPlutusCore.Core.Type

import Control.Lens

-- | Get all the direct child 'name a's of the given 'Term' from 'LamAbs'es.
termBinds :: Traversal' (Term name uni fun ann) name
termBinds :: (name -> f name)
-> Term name uni fun ann -> f (Term name uni fun ann)
termBinds name -> f name
f = \case
    LamAbs ann
ann name
n Term name uni fun ann
t -> name -> f name
f name
n f name
-> (name -> Term name uni fun ann) -> f (Term name uni fun ann)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \name
n' -> 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
n' Term name uni fun ann
t
    Term name uni fun ann
x              -> Term name uni fun ann -> f (Term name uni fun ann)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Term name uni fun ann
x

-- | Get all the direct child 'name a's of the given 'Term' from 'Var's.
termVars :: Traversal' (Term name uni fun ann) name
termVars :: (name -> f name)
-> Term name uni fun ann -> f (Term name uni fun ann)
termVars name -> f name
f = \case
    Var ann
ann name
n -> ann -> name -> Term name uni fun ann
forall name (uni :: * -> *) fun ann.
ann -> name -> Term name uni fun ann
Var ann
ann (name -> Term name uni fun ann)
-> f name -> f (Term name uni fun ann)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> name -> f name
f name
n
    Term name uni fun ann
x         -> Term name uni fun ann -> f (Term name uni fun ann)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Term name uni fun ann
x

-- | Get all the direct child 'Unique's of the given 'Term'.
termUniques :: HasUniques (Term name uni fun ann) => Traversal' (Term name uni fun ann) Unique
termUniques :: Traversal' (Term name uni fun ann) Unique
termUniques Unique -> f Unique
f = \case
    LamAbs ann
ann name
n Term name uni fun ann
t -> (Unique -> f Unique) -> name -> f name
forall name unique. HasUnique name unique => Lens' name Unique
theUnique Unique -> f Unique
f name
n f name
-> (name -> Term name uni fun ann) -> f (Term name uni fun ann)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \name
n' -> 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
n' Term name uni fun ann
t
    Var ann
ann name
n      -> (Unique -> f Unique) -> name -> f name
forall name unique. HasUnique name unique => Lens' name Unique
theUnique Unique -> f Unique
f name
n f name
-> (name -> Term name uni fun ann) -> f (Term name uni fun ann)
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> ann -> name -> Term name uni fun ann
forall name (uni :: * -> *) fun ann.
ann -> name -> Term name uni fun ann
Var ann
ann
    Term name uni fun ann
x              -> Term name uni fun ann -> f (Term name uni fun ann)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Term name uni fun ann
x

{-# INLINE termSubterms #-}
-- | Get all the direct child 'Term's of the given 'Term'.
termSubterms :: Traversal' (Term name uni fun ann) (Term name uni fun ann)
termSubterms :: (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 -> f (Term name uni fun ann)
f = \case
    LamAbs ann
ann name
n Term name uni fun ann
t  -> 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
n (Term name uni fun ann -> Term name uni fun ann)
-> f (Term name uni fun ann) -> f (Term name uni fun ann)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Term name uni fun ann -> f (Term name uni fun ann)
f Term name uni fun ann
t
    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 -> Term name uni fun ann)
-> f (Term name uni fun ann)
-> f (Term name uni fun ann -> Term name uni fun ann)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Term name uni fun ann -> f (Term name uni fun ann)
f Term name uni fun ann
t1 f (Term name uni fun ann -> Term name uni fun ann)
-> f (Term name uni fun ann) -> f (Term name uni fun ann)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Term name uni fun ann -> f (Term name uni fun ann)
f 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)
-> f (Term name uni fun ann) -> f (Term name uni fun ann)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Term name uni fun ann -> f (Term name uni fun ann)
f 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)
-> f (Term name uni fun ann) -> f (Term name uni fun ann)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Term name uni fun ann -> f (Term name uni fun ann)
f Term name uni fun ann
t
    e :: Term name uni fun ann
e@Error {}      -> Term name uni fun ann -> f (Term name uni fun ann)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Term name uni fun ann
e
    v :: Term name uni fun ann
v@Var {}        -> Term name uni fun ann -> f (Term name uni fun ann)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Term name uni fun ann
v
    c :: Term name uni fun ann
c@Constant {}   -> Term name uni fun ann -> f (Term name uni fun ann)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Term name uni fun ann
c
    b :: Term name uni fun ann
b@Builtin {}    -> Term name uni fun ann -> f (Term name uni fun ann)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Term name uni fun ann
b

-- | Get all the transitive child 'Term's of the given 'Term'.
termSubtermsDeep :: Fold (Term name uni fun ann) (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)
termSubtermsDeep = ((Term name uni fun ann -> f (Term name uni fun ann))
 -> Term name uni fun ann -> f (Term name uni fun ann))
-> (Term name uni fun ann -> f (Term name uni fun ann))
-> Term name uni fun ann
-> f (Term name uni fun ann)
forall (f :: * -> *) a.
(Applicative f, Contravariant f) =>
LensLike' f a a -> LensLike' f a a
cosmosOf (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.
Traversal' (Term name uni fun ann) (Term name uni fun ann)
termSubterms

-- | Get all the transitive child 'Unique's of the given 'Term'.
termUniquesDeep :: HasUniques (Term name uni fun ann) => Fold (Term name uni fun ann) Unique
termUniquesDeep :: Fold (Term name uni fun ann) Unique
termUniquesDeep = (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.
Fold (Term name uni fun ann) (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))
-> ((Unique -> f Unique)
    -> Term name uni fun ann -> f (Term name uni fun ann))
-> (Unique -> f Unique)
-> Term name uni fun ann
-> f (Term name uni fun ann)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Unique -> f Unique)
-> Term name uni fun ann -> f (Term name uni fun ann)
forall name (uni :: * -> *) fun ann.
HasUniques (Term name uni fun ann) =>
Traversal' (Term name uni fun ann) Unique
termUniques