{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PatternSynonyms #-}

{-| Floats immediately-applied lambdas ("let bindings") outwards, as long as doing so
cannot cause any expression to be evaluated more than before.

This can unlock further optimizations, such as case-constr and force-delay cancellation.

Specifically, it floats bindings from the following positions: @case@ scrutinee,
@force@ body, and the function of @apply@.

If we don't care about the order of effects, we can also float bindings from @apply@
arguments and @constr@ arguments, but there's no evidence that doing so is beneficial,
so we don't do it.

This pass only floats one bindings from each location. Since the pass runs once per
simplifier iteration, this should be sufficient. Floating out multi-lets at once would
complicate things, since we'd need to consider both @(\a b. M) A B@ and
@(\a. (\b -> M) B) A@. -}
module UntypedPlutusCore.Transform.LetFloatOut (letFloatOut) where

import PlutusCore qualified as PLC
import UntypedPlutusCore.Core
import UntypedPlutusCore.Transform.Optimizer
  ( OptimizerT
  , recordOptimization
  , pattern LetFloatOutStage
  )

import Control.Lens (transformOf)

letFloatOut
  :: ( PLC.MonadQuote m
     , PLC.Rename (Term name uni fun a)
     )
  => Term name uni fun a
  -> OptimizerT name uni fun a m (Term name uni fun a)
letFloatOut :: forall (m :: * -> *) name (uni :: * -> *) fun a.
(MonadQuote m, Rename (Term name uni fun a)) =>
Term name uni fun a
-> OptimizerT name uni fun a m (Term name uni fun a)
letFloatOut Term name uni fun a
term = do
  Term name uni fun a
result <- ASetter
  (Term name uni fun a)
  (Term name uni fun a)
  (Term name uni fun a)
  (Term name uni fun a)
-> (Term name uni fun a -> Term name uni fun a)
-> Term name uni fun a
-> Term name uni fun a
forall a b. ASetter a b a b -> (b -> b) -> a -> b
transformOf ASetter
  (Term name uni fun a)
  (Term name uni fun a)
  (Term name uni fun a)
  (Term name uni fun a)
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 a -> Term name uni fun a
forall name (uni :: * -> *) fun a.
Term name uni fun a -> Term name uni fun a
processTerm (Term name uni fun a -> Term name uni fun a)
-> OptimizerT name uni fun a m (Term name uni fun a)
-> OptimizerT name uni fun a m (Term name uni fun a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Term name uni fun a
-> OptimizerT name uni fun a m (Term name uni fun a)
forall a (m :: * -> *). (Rename a, MonadQuote m) => a -> m a
forall (m :: * -> *).
MonadQuote m =>
Term name uni fun a -> m (Term name uni fun a)
PLC.rename Term name uni fun a
term
  Term name uni fun a
-> OptStage
-> Term name uni fun a
-> OptimizerT name uni fun a m ()
forall (m :: * -> *) name (uni :: * -> *) fun a.
Monad m =>
Term name uni fun a
-> OptStage
-> Term name uni fun a
-> OptimizerT name uni fun a m ()
recordOptimization Term name uni fun a
term OptStage
LetFloatOutStage Term name uni fun a
result
  Term name uni fun a
-> OptimizerT name uni fun a m (Term name uni fun a)
forall a. a -> OptimizerT name uni fun a m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Term name uni fun a
result

processTerm :: Term name uni fun a -> Term name uni fun a
processTerm :: forall name (uni :: * -> *) fun a.
Term name uni fun a -> Term name uni fun a
processTerm = \case
  Case a
ca (Apply a
aa (LamAbs a
la name
x Term name uni fun a
body) Term name uni fun a
rhs) Vector (Term name uni fun a)
branches ->
    a
-> Term name uni fun a
-> Term name uni fun a
-> Term name uni fun a
forall name (uni :: * -> *) fun ann.
ann
-> Term name uni fun ann
-> Term name uni fun ann
-> Term name uni fun ann
Apply a
aa (a -> name -> Term name uni fun a -> Term name uni fun a
forall name (uni :: * -> *) fun ann.
ann -> name -> Term name uni fun ann -> Term name uni fun ann
LamAbs a
la name
x (a
-> Term name uni fun a
-> Vector (Term name uni fun a)
-> Term name uni fun a
forall name (uni :: * -> *) fun ann.
ann
-> Term name uni fun ann
-> Vector (Term name uni fun ann)
-> Term name uni fun ann
Case a
ca Term name uni fun a
body Vector (Term name uni fun a)
branches)) Term name uni fun a
rhs
  Force a
fa (Apply a
aa (LamAbs a
la name
x Term name uni fun a
body) Term name uni fun a
rhs) ->
    a
-> Term name uni fun a
-> Term name uni fun a
-> Term name uni fun a
forall name (uni :: * -> *) fun ann.
ann
-> Term name uni fun ann
-> Term name uni fun ann
-> Term name uni fun ann
Apply a
aa (a -> name -> Term name uni fun a -> Term name uni fun a
forall name (uni :: * -> *) fun ann.
ann -> name -> Term name uni fun ann -> Term name uni fun ann
LamAbs a
la name
x (a -> Term name uni fun a -> Term name uni fun a
forall name (uni :: * -> *) fun ann.
ann -> Term name uni fun ann -> Term name uni fun ann
Force a
fa Term name uni fun a
body)) Term name uni fun a
rhs
  Apply a
aa (Apply a
aa' (LamAbs a
la name
x Term name uni fun a
body) Term name uni fun a
rhs) Term name uni fun a
arg ->
    a
-> Term name uni fun a
-> Term name uni fun a
-> Term name uni fun a
forall name (uni :: * -> *) fun ann.
ann
-> Term name uni fun ann
-> Term name uni fun ann
-> Term name uni fun ann
Apply a
aa' (a -> name -> Term name uni fun a -> Term name uni fun a
forall name (uni :: * -> *) fun ann.
ann -> name -> Term name uni fun ann -> Term name uni fun ann
LamAbs a
la name
x (a
-> Term name uni fun a
-> Term name uni fun a
-> Term name uni fun a
forall name (uni :: * -> *) fun ann.
ann
-> Term name uni fun ann
-> Term name uni fun ann
-> Term name uni fun ann
Apply a
aa Term name uni fun a
body Term name uni fun a
arg)) Term name uni fun a
rhs
  Term name uni fun a
other -> Term name uni fun a
other