{-# LANGUAGE GADTs            #-}
{-# LANGUAGE RankNTypes       #-}
{-# LANGUAGE TypeApplications #-}

module UntypedPlutusCore.Simplify (
    module Opts,
    simplifyTerm,
    simplifyProgram,
    simplifyProgramWithTrace,
    InlineHints (..),
    termSimplifier,
    module UntypedPlutusCore.Transform.Simplifier,
) where

import PlutusCore.Compiler.Types
import PlutusCore.Default qualified as PLC
import PlutusCore.Default.Builtins
import PlutusCore.Name.Unique
import UntypedPlutusCore.Core.Type
import UntypedPlutusCore.Simplify.Opts as Opts
import UntypedPlutusCore.Transform.CaseOfCase
import UntypedPlutusCore.Transform.CaseReduce
import UntypedPlutusCore.Transform.Cse
import UntypedPlutusCore.Transform.FloatDelay (floatDelay)
import UntypedPlutusCore.Transform.ForceCaseDelay (forceCaseDelay)
import UntypedPlutusCore.Transform.ForceDelay (forceDelay)
import UntypedPlutusCore.Transform.Inline (InlineHints (..), inline)
import UntypedPlutusCore.Transform.Simplifier

import Control.Monad
import Data.List as List (foldl')
import Data.Typeable
import Data.Vector.Orphans ()

simplifyProgram ::
    forall name uni fun m a.
    (Compiling m uni fun name a) =>
    SimplifyOpts name a ->
    BuiltinSemanticsVariant fun ->
    Program name uni fun a ->
    m (Program name uni fun a)
simplifyProgram :: forall name (uni :: * -> *) fun (m :: * -> *) a.
Compiling m uni fun name a =>
SimplifyOpts name a
-> BuiltinSemanticsVariant fun
-> Program name uni fun a
-> m (Program name uni fun a)
simplifyProgram SimplifyOpts name a
opts BuiltinSemanticsVariant fun
builtinSemanticsVariant (Program a
a Version
v Term name uni fun a
t) =
  a -> Version -> Term name uni fun a -> Program name uni fun a
forall name (uni :: * -> *) fun ann.
ann -> Version -> Term name uni fun ann -> Program name uni fun ann
Program a
a Version
v (Term name uni fun a -> Program name uni fun a)
-> m (Term name uni fun a) -> m (Program name uni fun a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> SimplifyOpts name a
-> BuiltinSemanticsVariant fun
-> Term name uni fun a
-> m (Term name uni fun a)
forall name (uni :: * -> *) fun (m :: * -> *) a.
Compiling m uni fun name a =>
SimplifyOpts name a
-> BuiltinSemanticsVariant fun
-> Term name uni fun a
-> m (Term name uni fun a)
simplifyTerm SimplifyOpts name a
opts BuiltinSemanticsVariant fun
builtinSemanticsVariant Term name uni fun a
t

simplifyProgramWithTrace ::
    forall name uni fun m a.
    (Compiling m uni fun name a) =>
    SimplifyOpts name a ->
    BuiltinSemanticsVariant fun ->
    Program name uni fun a ->
    m (Program name uni fun a, SimplifierTrace name uni fun a)
simplifyProgramWithTrace :: forall name (uni :: * -> *) fun (m :: * -> *) a.
Compiling m uni fun name a =>
SimplifyOpts name a
-> BuiltinSemanticsVariant fun
-> Program name uni fun a
-> m (Program name uni fun a, SimplifierTrace name uni fun a)
simplifyProgramWithTrace SimplifyOpts name a
opts BuiltinSemanticsVariant fun
builtinSemanticsVariant (Program a
a Version
v Term name uni fun a
t) = do
  (Term name uni fun a
result, SimplifierTrace name uni fun a
trace) <-
    SimplifierT name uni fun a m (Term name uni fun a)
-> m (Term name uni fun a, SimplifierTrace name uni fun a)
forall name (uni :: * -> *) fun ann (m :: * -> *) a.
SimplifierT name uni fun ann m a
-> m (a, SimplifierTrace name uni fun ann)
runSimplifierT
      (SimplifierT name uni fun a m (Term name uni fun a)
 -> m (Term name uni fun a, SimplifierTrace name uni fun a))
-> SimplifierT name uni fun a m (Term name uni fun a)
-> m (Term name uni fun a, SimplifierTrace name uni fun a)
forall a b. (a -> b) -> a -> b
$ SimplifyOpts name a
-> BuiltinSemanticsVariant fun
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
forall name (uni :: * -> *) fun (m :: * -> *) a.
Compiling m uni fun name a =>
SimplifyOpts name a
-> BuiltinSemanticsVariant fun
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
termSimplifier SimplifyOpts name a
opts BuiltinSemanticsVariant fun
builtinSemanticsVariant Term name uni fun a
t
  (Program name uni fun a, SimplifierTrace name uni fun a)
-> m (Program name uni fun a, SimplifierTrace name uni fun a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> Version -> Term name uni fun a -> Program name uni fun a
forall name (uni :: * -> *) fun ann.
ann -> Version -> Term name uni fun ann -> Program name uni fun ann
Program a
a Version
v Term name uni fun a
result, SimplifierTrace name uni fun a
trace)

simplifyTerm ::
    forall name uni fun m a.
    (Compiling m uni fun name a) =>
    SimplifyOpts name a ->
    BuiltinSemanticsVariant fun ->
    Term name uni fun a ->
    m (Term name uni fun a)
simplifyTerm :: forall name (uni :: * -> *) fun (m :: * -> *) a.
Compiling m uni fun name a =>
SimplifyOpts name a
-> BuiltinSemanticsVariant fun
-> Term name uni fun a
-> m (Term name uni fun a)
simplifyTerm SimplifyOpts name a
opts BuiltinSemanticsVariant fun
builtinSemanticsVariant Term name uni fun a
term =
  SimplifierT name uni fun a m (Term name uni fun a)
-> m (Term name uni fun a)
forall (m :: * -> *) name (uni :: * -> *) fun ann a.
Monad m =>
SimplifierT name uni fun ann m a -> m a
evalSimplifierT (SimplifierT name uni fun a m (Term name uni fun a)
 -> m (Term name uni fun a))
-> SimplifierT name uni fun a m (Term name uni fun a)
-> m (Term name uni fun a)
forall a b. (a -> b) -> a -> b
$ SimplifyOpts name a
-> BuiltinSemanticsVariant fun
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
forall name (uni :: * -> *) fun (m :: * -> *) a.
Compiling m uni fun name a =>
SimplifyOpts name a
-> BuiltinSemanticsVariant fun
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
termSimplifier SimplifyOpts name a
opts BuiltinSemanticsVariant fun
builtinSemanticsVariant Term name uni fun a
term

termSimplifier ::
    forall name uni fun m a.
    (Compiling m uni fun name a) =>
    SimplifyOpts name a ->
    BuiltinSemanticsVariant fun ->
    Term name uni fun a ->
    SimplifierT name uni fun a m (Term name uni fun a)
termSimplifier :: forall name (uni :: * -> *) fun (m :: * -> *) a.
Compiling m uni fun name a =>
SimplifyOpts name a
-> BuiltinSemanticsVariant fun
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
termSimplifier SimplifyOpts name a
opts BuiltinSemanticsVariant fun
builtinSemanticsVariant =
    Int
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
simplifyNTimes (SimplifyOpts name a -> Int
forall name a. SimplifyOpts name a -> Int
_soMaxSimplifierIterations SimplifyOpts name a
opts) (Term name uni fun a
 -> SimplifierT name uni fun a m (Term name uni fun a))
-> (Term name uni fun a
    -> SimplifierT name uni fun a m (Term name uni fun a))
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> Int
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
cseNTimes Int
cseTimes
  where
    -- Run the simplifier @n@ times
    simplifyNTimes ::
      Int ->
      Term name uni fun a ->
      SimplifierT name uni fun a m (Term name uni fun a)
    simplifyNTimes :: Int
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
simplifyNTimes Int
n = ((Term name uni fun a
  -> SimplifierT name uni fun a m (Term name uni fun a))
 -> (Term name uni fun a
     -> SimplifierT name uni fun a m (Term name uni fun a))
 -> Term name uni fun a
 -> SimplifierT name uni fun a m (Term name uni fun a))
-> (Term name uni fun a
    -> SimplifierT name uni fun a m (Term name uni fun a))
-> [Term name uni fun a
    -> SimplifierT name uni fun a m (Term name uni fun a)]
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
List.foldl' (Term name uni fun a
 -> SimplifierT name uni fun a m (Term name uni fun a))
-> (Term name uni fun a
    -> SimplifierT name uni fun a m (Term name uni fun a))
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
(>=>) Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
forall a. a -> SimplifierT name uni fun a m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Term name uni fun a
  -> SimplifierT name uni fun a m (Term name uni fun a)]
 -> Term name uni fun a
 -> SimplifierT name uni fun a m (Term name uni fun a))
-> [Term name uni fun a
    -> SimplifierT name uni fun a m (Term name uni fun a)]
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
forall a b. (a -> b) -> a -> b
$ (Int
 -> Term name uni fun a
 -> SimplifierT name uni fun a m (Term name uni fun a))
-> [Int]
-> [Term name uni fun a
    -> SimplifierT name uni fun a m (Term name uni fun a)]
forall a b. (a -> b) -> [a] -> [b]
map Int
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
simplifyStep [Int
1..Int
n]

    -- Run CSE @n@ times, interleaved with the simplifier.
    -- See Note [CSE]
    cseNTimes ::
      Int ->
      Term name uni fun a ->
      SimplifierT name uni fun a m (Term name uni fun a)
    cseNTimes :: Int
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
cseNTimes Int
n = ((Term name uni fun a
  -> SimplifierT name uni fun a m (Term name uni fun a))
 -> (Term name uni fun a
     -> SimplifierT name uni fun a m (Term name uni fun a))
 -> Term name uni fun a
 -> SimplifierT name uni fun a m (Term name uni fun a))
-> (Term name uni fun a
    -> SimplifierT name uni fun a m (Term name uni fun a))
-> [Term name uni fun a
    -> SimplifierT name uni fun a m (Term name uni fun a)]
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (Term name uni fun a
 -> SimplifierT name uni fun a m (Term name uni fun a))
-> (Term name uni fun a
    -> SimplifierT name uni fun a m (Term name uni fun a))
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
(>=>) Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
forall a. a -> SimplifierT name uni fun a m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Term name uni fun a
  -> SimplifierT name uni fun a m (Term name uni fun a)]
 -> Term name uni fun a
 -> SimplifierT name uni fun a m (Term name uni fun a))
-> [Term name uni fun a
    -> SimplifierT name uni fun a m (Term name uni fun a)]
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
forall a b. (a -> b) -> a -> b
$ (Int
 -> [Term name uni fun a
     -> SimplifierT name uni fun a m (Term name uni fun a)])
-> [Int]
-> [Term name uni fun a
    -> SimplifierT name uni fun a m (Term name uni fun a)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\Int
i -> [Int
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
cseStep Int
i, Int
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
simplifyStep Int
i]) [Int
1..Int
n]

    -- generate simplification step
    simplifyStep ::
      Int ->
      Term name uni fun a ->
      SimplifierT name uni fun a m (Term name uni fun a)
    simplifyStep :: Int
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
simplifyStep Int
_ =
        Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
forall (m :: * -> *) name (uni :: * -> *) fun a.
(MonadQuote m, Rename (Term name uni fun a),
 HasUnique name TermUnique) =>
Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
floatDelay
        (Term name uni fun a
 -> SimplifierT name uni fun a m (Term name uni fun a))
-> (Term name uni fun a
    -> SimplifierT name uni fun a m (Term name uni fun a))
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
forall (m :: * -> *) name (uni :: * -> *) fun a.
Monad m =>
Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
forceCaseDelay
        (Term name uni fun a
 -> SimplifierT name uni fun a m (Term name uni fun a))
-> (Term name uni fun a
    -> SimplifierT name uni fun a m (Term name uni fun a))
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> case (forall {k} (a :: k) (b :: k).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
forall (a :: * -> *) (b :: * -> *).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
eqT @uni @PLC.DefaultUni, forall {k} (a :: k) (b :: k).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
forall a b. (Typeable a, Typeable b) => Maybe (a :~: b)
eqT @fun @DefaultFun) of
            (Just uni :~: DefaultUni
Refl, Just fun :~: DefaultFun
Refl) -> BuiltinSemanticsVariant fun
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
forall (uni :: * -> *) fun (m :: * -> *) name a.
(uni ~ DefaultUni, fun ~ DefaultFun, Monad m) =>
BuiltinSemanticsVariant fun
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
forceDelay BuiltinSemanticsVariant fun
builtinSemanticsVariant
            (Maybe (uni :~: DefaultUni), Maybe (fun :~: DefaultFun))
_                      -> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
forall a. a -> SimplifierT name uni fun a m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
        (Term name uni fun a
 -> SimplifierT name uni fun a m (Term name uni fun a))
-> (Term name uni fun a
    -> SimplifierT name uni fun a m (Term name uni fun a))
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
caseOfCase'
        (Term name uni fun a
 -> SimplifierT name uni fun a m (Term name uni fun a))
-> (Term name uni fun a
    -> SimplifierT name uni fun a m (Term name uni fun a))
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
forall (m :: * -> *) (uni :: * -> *) name fun a.
(Monad m, CaseBuiltin uni) =>
Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
caseReduce
        (Term name uni fun a
 -> SimplifierT name uni fun a m (Term name uni fun a))
-> (Term name uni fun a
    -> SimplifierT name uni fun a m (Term name uni fun a))
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> Size
-> Bool
-> Bool
-> InlineHints name a
-> BuiltinSemanticsVariant fun
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
forall name (uni :: * -> *) fun (m :: * -> *) a.
ExternalConstraints name uni fun m =>
Size
-> Bool
-> Bool
-> InlineHints name a
-> BuiltinSemanticsVariant fun
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
inline
              (SimplifyOpts name a -> Size
forall name a. SimplifyOpts name a -> Size
_soInlineCallsiteGrowth SimplifyOpts name a
opts)
              (SimplifyOpts name a -> Bool
forall name a. SimplifyOpts name a -> Bool
_soInlineConstants SimplifyOpts name a
opts)
              (SimplifyOpts name a -> Bool
forall name a. SimplifyOpts name a -> Bool
_soPreserveLogging SimplifyOpts name a
opts)
              (SimplifyOpts name a -> InlineHints name a
forall name a. SimplifyOpts name a -> InlineHints name a
_soInlineHints SimplifyOpts name a
opts)
              BuiltinSemanticsVariant fun
builtinSemanticsVariant

    caseOfCase' ::
      Term name uni fun a ->
      SimplifierT name uni fun a m (Term name uni fun a)
    caseOfCase' :: Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
caseOfCase' = case forall {k} (a :: k) (b :: k).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
forall a b. (Typeable a, Typeable b) => Maybe (a :~: b)
eqT @fun @DefaultFun of
      Just fun :~: DefaultFun
Refl -> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
forall fun (m :: * -> *) (uni :: * -> *) name a.
(fun ~ DefaultFun, Monad m, CaseBuiltin uni, GEq uni, Closed uni,
 Everywhere uni Eq) =>
Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
caseOfCase
      Maybe (fun :~: DefaultFun)
Nothing   -> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
forall a. a -> SimplifierT name uni fun a m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

    cseStep ::
      Int ->
      Term name uni fun a ->
      SimplifierT name uni fun a m (Term name uni fun a)
    cseStep :: Int
-> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
cseStep Int
_ =
      case (forall {k} (a :: k) (b :: k).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
forall a b. (Typeable a, Typeable b) => Maybe (a :~: b)
eqT @name @Name, forall {k} (a :: k) (b :: k).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
forall (a :: * -> *) (b :: * -> *).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
eqT @uni @PLC.DefaultUni) of
        (Just name :~: Name
Refl, Just uni :~: DefaultUni
Refl) -> BuiltinSemanticsVariant fun
-> Term Name uni fun a
-> SimplifierT Name uni fun a m (Term Name uni fun a)
forall (m :: * -> *) (uni :: * -> *) fun ann.
(MonadQuote m, Hashable (Term Name uni fun ()),
 Rename (Term Name uni fun ann), ToBuiltinMeaning uni fun) =>
BuiltinSemanticsVariant fun
-> Term Name uni fun ann
-> SimplifierT Name uni fun ann m (Term Name uni fun ann)
cse BuiltinSemanticsVariant fun
builtinSemanticsVariant
        (Maybe (name :~: Name), Maybe (uni :~: DefaultUni))
_                      -> Term name uni fun a
-> SimplifierT name uni fun a m (Term name uni fun a)
forall a. a -> SimplifierT name uni fun a m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

    cseTimes :: Int
cseTimes = if SimplifyOpts name a -> Bool
forall name a. SimplifyOpts name a -> Bool
_soConservativeOpts SimplifyOpts name a
opts then Int
0 else SimplifyOpts name a -> Int
forall name a. SimplifyOpts name a -> Int
_soMaxCseIterations SimplifyOpts name a
opts