{- |
Unrolling a recursive function involves explicitly laying out some or all of the recursive steps, rather than relying on recursion via a fixed-point combinator.
In UPLC, the typical `pfix` implementation uses a Y-combinator under the hood. Each recursive step incurs additional evaluation costs (CPU and memory) due to
the execution of the Y-combinator. By eliminating these costs in each recursive step, unrolled functions reduce execution overhead. However, since each recursive
step is explicitly laid out, unrolled functions consume more script size.

There are various unrolling strategies available. It is important to carefully study the implications of each strategy, as they may impose different requirements,
such as hard limit on recursion depth.

Generally for all strategies, calling the "recursion step"--that is @r@ in @fix $ \r a b ... -> ...@--multiple times will result in exponential increase of the script
size in each step of unrolling. This is because everytime the recursion step is called, it generates UPLC code for that specific branch(and all the subsequent branches
down the unrolling!). To prevent this, it is possible to have @plet r $ \r' -> <..r' can be used multiple times..>@.

Also, due to some of the other optimization performed; specifically @pmatch@'s branch detection which requires hashing terms of all subsequent deconstruction branches,
only tail-call recursion would perform fast enough to be used for large unrolling.
-}
module Plutarch.Unroll (punrollBound, punrollBound', punrollUnbound, punrollUnboundWhole) where

import Plutarch.Internal.Fix (pfix)
import Plutarch.Internal.PLam (plam)
import Plutarch.Internal.Term (Term, (#$), (:-->))

{- |
This is simplified version of @punrollBound'@ without doing additional recursion on Haskell level value.

@since 1.10.0
-}
punrollBound ::
  forall a b s.
  Integer ->
  Term s (a :--> b) ->
  (Term s (a :--> b) -> Term s (a :--> b)) ->
  Term s (a :--> b)
punrollBound :: forall (a :: PType) (b :: PType) (s :: S).
Integer
-> Term s (a :--> b)
-> (Term s (a :--> b) -> Term s (a :--> b))
-> Term s (a :--> b)
punrollBound Integer
d Term s (a :--> b)
def Term s (a :--> b) -> Term s (a :--> b)
f = Integer
-> (() -> Term s (a :--> b))
-> ((() -> Term s (a :--> b)) -> () -> Term s (a :--> b))
-> ()
-> Term s (a :--> b)
forall (a :: PType) (b :: PType) c (s :: S).
Integer
-> (c -> Term s (a :--> b))
-> ((c -> Term s (a :--> b)) -> c -> Term s (a :--> b))
-> c
-> Term s (a :--> b)
punrollBound' Integer
d (Term s (a :--> b) -> () -> Term s (a :--> b)
forall a b. a -> b -> a
const Term s (a :--> b)
def) (\() -> Term s (a :--> b)
g () -> Term s (a :--> b) -> Term s (a :--> b)
f (() -> Term s (a :--> b)
g ())) ()

{- |
The first argument specifies the unrolling depth.
The second argument defines the fallback behavior when the recursion depth exceeds the provided unrolling depth.

The fixed-point implementation provided requires a Haskell-level value @c@ and a Plutarch function of type `Term s (a :--> b)`. The functional for the recursion is passed as a Haskell function.
The inclusion of the additional, arbitrary Haskell value (typed @c@) enables further optimization by allowing pre-computation of constant values that depend only on the recursion depth.

This function will be used in a very niche situations. Using Haskell-level value for constant replacement is only practical on a single branch recursion with constant value that needs to be added
on each step. @plength@ is one of the niche use case.

@since 1.10.0
-}
punrollBound' ::
  forall a b c s.
  Integer ->
  (c -> Term s (a :--> b)) ->
  ((c -> Term s (a :--> b)) -> c -> Term s (a :--> b)) ->
  c ->
  Term s (a :--> b)
punrollBound' :: forall (a :: PType) (b :: PType) c (s :: S).
Integer
-> (c -> Term s (a :--> b))
-> ((c -> Term s (a :--> b)) -> c -> Term s (a :--> b))
-> c
-> Term s (a :--> b)
punrollBound' Integer
0 c -> Term s (a :--> b)
def (c -> Term s (a :--> b)) -> c -> Term s (a :--> b)
_ c
c = c -> Term s (a :--> b)
def c
c
punrollBound' Integer
d c -> Term s (a :--> b)
def (c -> Term s (a :--> b)) -> c -> Term s (a :--> b)
f c
c = (c -> Term s (a :--> b)) -> c -> Term s (a :--> b)
f (Integer
-> (c -> Term s (a :--> b))
-> ((c -> Term s (a :--> b)) -> c -> Term s (a :--> b))
-> c
-> Term s (a :--> b)
forall (a :: PType) (b :: PType) c (s :: S).
Integer
-> (c -> Term s (a :--> b))
-> ((c -> Term s (a :--> b)) -> c -> Term s (a :--> b))
-> c
-> Term s (a :--> b)
punrollBound' (Integer
d Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1) c -> Term s (a :--> b)
def (c -> Term s (a :--> b)) -> c -> Term s (a :--> b)
f) c
c

{- |
Unroll given amount of steps, and for rest, uses `pfix` to support unbound recursion.

@since 1.10.0
-}
punrollUnbound :: forall a b s. Integer -> (Term s (a :--> b) -> Term s (a :--> b)) -> Term s (a :--> b)
punrollUnbound :: forall (a :: PType) (b :: PType) (s :: S).
Integer
-> (Term s (a :--> b) -> Term s (a :--> b)) -> Term s (a :--> b)
punrollUnbound Integer
0 Term s (a :--> b) -> Term s (a :--> b)
f = Term s (((a :--> b) :--> (a :--> b)) :--> (a :--> b))
forall (s :: S) (a :: PType) (b :: PType).
Term s (((a :--> b) :--> (a :--> b)) :--> (a :--> b))
pfix Term s (((a :--> b) :--> (a :--> b)) :--> (a :--> b))
-> Term s ((a :--> b) :--> (a :--> b)) -> Term s (a :--> b)
forall (s :: S) (a :: PType) (b :: PType).
Term s (a :--> b) -> Term s a -> Term s b
#$ (Term s (a :--> b) -> Term s (a :--> b))
-> Term s ((a :--> b) :--> (a :--> b))
forall a (b :: PType) (s :: S) (c :: PType).
(PLamN a b s, HasCallStack) =>
(Term s c -> a) -> Term s (c :--> b)
forall (c :: PType).
HasCallStack =>
(Term s c -> Term s (a :--> b)) -> Term s (c :--> (a :--> b))
plam Term s (a :--> b) -> Term s (a :--> b)
f
punrollUnbound Integer
d Term s (a :--> b) -> Term s (a :--> b)
f = Term s (a :--> b) -> Term s (a :--> b)
f (Integer
-> (Term s (a :--> b) -> Term s (a :--> b)) -> Term s (a :--> b)
forall (a :: PType) (b :: PType) (s :: S).
Integer
-> (Term s (a :--> b) -> Term s (a :--> b)) -> Term s (a :--> b)
punrollUnbound (Integer
d Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
1) Term s (a :--> b) -> Term s (a :--> b)
f)

{- |
Uses `pfix` to recurse unrolled function itself. Unlike @punrollUnbound@, this function uses unrolled instructions
within `pfix` recursions.

This should perform better than @punrollUnbound@ when a function requires a large recursion depth.

@since 1.10.0
-}
punrollUnboundWhole :: forall a b s. Integer -> (Term s (a :--> b) -> Term s (a :--> b)) -> Term s (a :--> b)
punrollUnboundWhole :: forall (a :: PType) (b :: PType) (s :: S).
Integer
-> (Term s (a :--> b) -> Term s (a :--> b)) -> Term s (a :--> b)
punrollUnboundWhole Integer
d Term s (a :--> b) -> Term s (a :--> b)
f = Term s (((a :--> b) :--> (a :--> b)) :--> (a :--> b))
forall (s :: S) (a :: PType) (b :: PType).
Term s (((a :--> b) :--> (a :--> b)) :--> (a :--> b))
pfix Term s (((a :--> b) :--> (a :--> b)) :--> (a :--> b))
-> Term s ((a :--> b) :--> (a :--> b)) -> Term s (a :--> b)
forall (s :: S) (a :: PType) (b :: PType).
Term s (a :--> b) -> Term s a -> Term s b
#$ (Term s (a :--> b) -> Term s (a :--> b))
-> Term s ((a :--> b) :--> (a :--> b))
forall a (b :: PType) (s :: S) (c :: PType).
(PLamN a b s, HasCallStack) =>
(Term s c -> a) -> Term s (c :--> b)
forall (c :: PType).
HasCallStack =>
(Term s c -> Term s (a :--> b)) -> Term s (c :--> (a :--> b))
plam ((Term s (a :--> b) -> Term s (a :--> b))
 -> Term s ((a :--> b) :--> (a :--> b)))
-> (Term s (a :--> b) -> Term s (a :--> b))
-> Term s ((a :--> b) :--> (a :--> b))
forall a b. (a -> b) -> a -> b
$ \Term s (a :--> b)
r -> Integer
-> Term s (a :--> b)
-> (Term s (a :--> b) -> Term s (a :--> b))
-> Term s (a :--> b)
forall (a :: PType) (b :: PType) (s :: S).
Integer
-> Term s (a :--> b)
-> (Term s (a :--> b) -> Term s (a :--> b))
-> Term s (a :--> b)
punrollBound Integer
d Term s (a :--> b)
r Term s (a :--> b) -> Term s (a :--> b)
f