{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

{-| Hoist fully-forced polymorphic builtins.

Before: @...force (force fstPair)...@
After: @(\fstPairForced -> ...fstPairForced...) (force (force fstPair))@ -}
module UntypedPlutusCore.Transform.PolyBuiltin
  ( polyBuiltin
  ) where

import PlutusCore (Name, getAnn)
import PlutusCore.Arity (Param (..), builtinArity)
import PlutusCore.Builtin (ToBuiltinMeaning (BuiltinSemanticsVariant))
import PlutusCore.Quote (MonadQuote, freshName)
import UntypedPlutusCore.Core.Plated (termSubterms, termSubtermsDeep)
import UntypedPlutusCore.Core.Type (Term (..))
import UntypedPlutusCore.Transform.Optimizer
  ( OptimizerT
  , recordOptimization
  , pattern PolyBuiltinStage
  )

import Control.Lens (foldlOf', transformOf)
import Data.Foldable qualified as F
import Data.HashMap.Strict (HashMap)
import Data.HashMap.Strict qualified as Map
import Data.Hashable (Hashable)
import Data.Proxy (Proxy (..))
import Data.Traversable (for)

polyBuiltin
  :: forall m uni fun a
   . ( MonadQuote m
     , Hashable fun
     , ToBuiltinMeaning uni fun
     )
  => BuiltinSemanticsVariant fun
  -> Term Name uni fun a
  -> OptimizerT Name uni fun a m (Term Name uni fun a)
polyBuiltin :: forall (m :: * -> *) (uni :: * -> *) fun a.
(MonadQuote m, Hashable fun, ToBuiltinMeaning uni fun) =>
BuiltinSemanticsVariant fun
-> Term Name uni fun a
-> OptimizerT Name uni fun a m (Term Name uni fun a)
polyBuiltin BuiltinSemanticsVariant fun
semvar Term Name uni fun a
term = do
  let numForces :: fun -> Int
      numForces :: fun -> Int
numForces = [Param] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Param] -> Int) -> (fun -> [Param]) -> fun -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Param -> Bool) -> [Param] -> [Param]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Param -> Param -> Bool
forall a. Eq a => a -> a -> Bool
== Param
TypeParam) ([Param] -> [Param]) -> (fun -> [Param]) -> fun -> [Param]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Proxy uni -> BuiltinSemanticsVariant fun -> fun -> [Param]
forall (uni :: * -> *) fun.
ToBuiltinMeaning uni fun =>
Proxy uni -> BuiltinSemanticsVariant fun -> fun -> [Param]
builtinArity (forall {k} (t :: k). Proxy t
forall (t :: * -> *). Proxy t
Proxy @uni) BuiltinSemanticsVariant fun
semvar

      -- If it is a fully forced polymorphic builtin, return `Just (ann, fun, number of forces)`.
      checkFullyForced :: Term Name uni fun a -> Maybe (a, fun, Int)
      checkFullyForced :: Term Name uni fun a -> Maybe (a, fun, Int)
checkFullyForced Term Name uni fun a
t = case Term Name uni fun a
t' of
        Builtin a
_ fun
fun
          | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
1
          , Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== fun -> Int
numForces fun
fun ->
              (a, fun, Int) -> Maybe (a, fun, Int)
forall a. a -> Maybe a
Just (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, fun
fun, Int
n)
        Term Name uni fun a
_ -> Maybe (a, fun, Int)
forall a. Maybe a
Nothing
        where
          (Int
n, Term Name uni fun a
t') = Term Name uni fun a -> (Int, Term Name uni fun a)
forall name (uni :: * -> *) fun a.
Term name uni fun a -> (Int, Term name uni fun a)
peelForces Term Name uni fun a
t

      candidates :: HashMap fun Int
      candidates :: HashMap fun Int
candidates =
        Getting
  (Endo (Endo (HashMap fun Int)))
  (Term Name uni fun a)
  (Term Name uni fun a)
-> (HashMap fun Int -> Term Name uni fun a -> HashMap fun Int)
-> HashMap fun Int
-> Term Name uni fun a
-> HashMap fun Int
forall r s a.
Getting (Endo (Endo r)) s a -> (r -> a -> r) -> r -> s -> r
foldlOf' Getting
  (Endo (Endo (HashMap fun Int)))
  (Term Name uni fun a)
  (Term Name uni fun a)
forall name (uni :: * -> *) fun ann (f :: * -> *).
(Contravariant 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)
termSubtermsDeep HashMap fun Int -> Term Name uni fun a -> HashMap fun Int
alg HashMap fun Int
forall k v. HashMap k v
Map.empty Term Name uni fun a
term
        where
          alg :: HashMap fun Int -> Term Name uni fun a -> HashMap fun Int
alg !HashMap fun Int
acc Term Name uni fun a
t = case Term Name uni fun a -> Maybe (a, fun, Int)
checkFullyForced Term Name uni fun a
t of
            Just (a
_, fun
fun, Int
n) -> fun -> Int -> HashMap fun Int -> HashMap fun Int
forall k v. Hashable k => k -> v -> HashMap k v -> HashMap k v
Map.insert fun
fun Int
n HashMap fun Int
acc
            Maybe (a, fun, Int)
Nothing -> HashMap fun Int
acc

  Term Name uni fun a
result <-
    if HashMap fun Int -> Bool
forall k v. HashMap k v -> Bool
Map.null HashMap fun Int
candidates
      then 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
term
      else do
        [(fun, (Name, Int))]
named <- [(fun, Int)]
-> ((fun, Int) -> OptimizerT Name uni fun a m (fun, (Name, Int)))
-> OptimizerT Name uni fun a m [(fun, (Name, Int))]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for (HashMap fun Int -> [(fun, Int)]
forall k v. HashMap k v -> [(k, v)]
Map.toList HashMap fun Int
candidates) (((fun, Int) -> OptimizerT Name uni fun a m (fun, (Name, Int)))
 -> OptimizerT Name uni fun a m [(fun, (Name, Int))])
-> ((fun, Int) -> OptimizerT Name uni fun a m (fun, (Name, Int)))
-> OptimizerT Name uni fun a m [(fun, (Name, Int))]
forall a b. (a -> b) -> a -> b
$ \(fun
fun, Int
n) -> do
          Name
name <- Text -> OptimizerT Name uni fun a m Name
forall (m :: * -> *). MonadQuote m => Text -> m Name
freshName Text
"forced"
          (fun, (Name, Int))
-> OptimizerT Name uni fun a m (fun, (Name, Int))
forall a. a -> OptimizerT Name uni fun a m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (fun
fun, (Name
name, Int
n))
        let binders :: HashMap fun (Name, Int)
binders = [(fun, (Name, Int))] -> HashMap fun (Name, Int)
forall k v. Hashable k => [(k, v)] -> HashMap k v
Map.fromList [(fun, (Name, Int))]
named
            subst :: Term Name uni fun a -> Term Name uni fun a
subst Term Name uni fun a
t = case Term Name uni fun a -> Maybe (a, fun, Int)
checkFullyForced Term Name uni fun a
t of
              Just (a
ann, fun
fun, Int
_)
                | Just (Name
name, Int
_) <- fun -> HashMap fun (Name, Int) -> Maybe (Name, Int)
forall k v. Hashable k => k -> HashMap k v -> Maybe v
Map.lookup fun
fun HashMap fun (Name, Int)
binders ->
                    a -> Name -> Term Name uni fun a
forall name (uni :: * -> *) fun ann.
ann -> name -> Term name uni fun ann
Var a
ann Name
name
              Maybe (a, fun, Int)
_ -> Term Name uni fun a
t
            rewritten :: Term Name uni fun a
rewritten = 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
subst Term Name uni fun a
term
        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
 -> OptimizerT Name uni fun a m (Term Name uni fun a))
-> Term Name uni fun a
-> OptimizerT Name uni fun a m (Term Name uni fun a)
forall a b. (a -> b) -> a -> b
$ a
-> [(fun, (Name, Int))]
-> Term Name uni fun a
-> Term Name uni fun a
forall (uni :: * -> *) fun a.
a
-> [(fun, (Name, Int))]
-> Term Name uni fun a
-> Term Name uni fun a
wrap (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
term) [(fun, (Name, Int))]
named Term Name uni fun a
rewritten

  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
PolyBuiltinStage 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

peelForces :: Term name uni fun a -> (Int, Term name uni fun a)
peelForces :: forall name (uni :: * -> *) fun a.
Term name uni fun a -> (Int, Term name uni fun a)
peelForces = Int -> Term name uni fun a -> (Int, Term name uni fun a)
forall {a} {name} {uni :: * -> *} {fun} {ann}.
Num a =>
a -> Term name uni fun ann -> (a, Term name uni fun ann)
go Int
0
  where
    go :: a -> Term name uni fun ann -> (a, Term name uni fun ann)
go !a
k (Force ann
_ Term name uni fun ann
t) = a -> Term name uni fun ann -> (a, Term name uni fun ann)
go (a
k a -> a -> a
forall a. Num a => a -> a -> a
+ a
1) Term name uni fun ann
t
    go !a
k Term name uni fun ann
t = (a
k, Term name uni fun ann
t)

wrap
  :: forall uni fun a
   . a
  -> [(fun, (Name, Int))]
  -> Term Name uni fun a
  -> Term Name uni fun a
wrap :: forall (uni :: * -> *) fun a.
a
-> [(fun, (Name, Int))]
-> Term Name uni fun a
-> Term Name uni fun a
wrap a
ann [(fun, (Name, Int))]
entries Term Name uni fun a
body =
  let lams :: Term Name uni fun a
lams = ((fun, (Name, Int)) -> Term Name uni fun a -> Term Name uni fun a)
-> Term Name uni fun a
-> [(fun, (Name, Int))]
-> Term Name uni fun a
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\(fun
_, (Name
name, Int
_)) -> 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
ann Name
name) Term Name uni fun a
body [(fun, (Name, Int))]
entries
   in (Term Name uni fun a -> (fun, (Name, Int)) -> Term Name uni fun a)
-> Term Name uni fun a
-> [(fun, (Name, Int))]
-> 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
F.foldl'
        (\Term Name uni fun a
acc (fun
fun, (Name
_, Int
n)) -> 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
ann Term Name uni fun a
acc (Int -> Term Name uni fun a -> Term Name uni fun a
applyForces Int
n (a -> fun -> Term Name uni fun a
forall name (uni :: * -> *) fun ann.
ann -> fun -> Term name uni fun ann
Builtin a
ann fun
fun)))
        Term Name uni fun a
lams
        [(fun, (Name, Int))]
entries
  where
    applyForces :: Int -> Term Name uni fun a -> Term Name uni fun a
    applyForces :: Int -> Term Name uni fun a -> Term Name uni fun a
applyForces Int
0 = Term Name uni fun a -> Term Name uni fun a
forall a. a -> a
id
    applyForces Int
n = Int -> Term Name uni fun a -> Term Name uni fun a
applyForces (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (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 b c a. (b -> c) -> (a -> b) -> a -> c
. 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
ann