{-# LANGUAGE PatternSynonyms #-}

{-|
Transform multi-argument applications into case-constr form.

This is essentially the reverse of 'CaseReduce'. An application of 3 or more
arguments, e.g., @f x y z@ is rewritten to @case (constr 0 [x, y, z]) [f]@,
which is cheaper, and smaller in AST size, though it is sometimes slightly
bigger in cbor size.

I cannot think of any case where this pass could enable further optimizations.
It can certainly destroy optimization opportunities, similar to CSE.
So it should only be run once, at the very end of the pipeline. CSE is also
run at the end, but it runs for multiple iterations interleaved with other
optimizations, because CSE can both destroy and enable further optimizations. -}
module UntypedPlutusCore.Transform.ApplyToCase (applyToCase) where

import Control.Lens (over)
import Data.Vector qualified as V
import PlutusPrelude (getAnn)
import UntypedPlutusCore.Core
import UntypedPlutusCore.Transform.Optimizer
  ( OptimizerT
  , recordOptimization
  , pattern ApplyToCaseStage
  )

minArgs :: Int
minArgs :: Int
minArgs = Int
3
{-# INLINE minArgs #-}

applyToCase
  :: Monad m
  => Term name uni fun a
  -> OptimizerT name uni fun a m (Term name uni fun a)
applyToCase :: forall (m :: * -> *) name (uni :: * -> *) fun a.
Monad m =>
Term name uni fun a
-> OptimizerT name uni fun a m (Term name uni fun a)
applyToCase Term name uni fun a
term = do
  let result :: Term name uni fun a
result = 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
  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
ApplyToCaseStage 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 Term name uni fun a
t = case Term name uni fun a
-> (Term name uni fun a, [(a, Term name uni fun a)])
forall name (uni :: * -> *) fun a.
Term name uni fun a
-> (Term name uni fun a, [(a, Term name uni fun a)])
splitApplication Term name uni fun a
t of
  (Term name uni fun a
fun, [(a, Term name uni fun a)]
args)
    | [(a, Term name uni fun a)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(a, Term name uni fun a)]
args Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
minArgs ->
        let ann :: a
ann = Term name uni fun a -> a
forall a. Term name uni fun a -> a
forall (f :: * -> *) a. HasAnn f => f a -> a
getAnn Term name uni fun a
t
         in 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
ann (a -> Word64 -> [Term name uni fun a] -> Term name uni fun a
forall name (uni :: * -> *) fun ann.
ann -> Word64 -> [Term name uni fun ann] -> Term name uni fun ann
Constr a
ann Word64
0 (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)
-> ((a, Term name uni fun a) -> Term name uni fun a)
-> (a, Term name uni fun a)
-> Term name uni fun a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a, Term name uni fun a) -> Term name uni fun a
forall a b. (a, b) -> b
snd ((a, Term name uni fun a) -> Term name uni fun a)
-> [(a, Term name uni fun a)] -> [Term name uni fun a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(a, Term name uni fun a)]
args)) (Term name uni fun a -> Vector (Term name uni fun a)
forall a. a -> Vector a
V.singleton (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
fun))
  (Term name uni fun a, [(a, Term name uni fun a)])
_ -> 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 s t a b. ASetter s t a b -> (a -> b) -> s -> t
over 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
t