{-# LANGUAGE ConstraintKinds  #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE GADTs            #-}
{-# LANGUAGE LambdaCase       #-}
{-# LANGUAGE TemplateHaskell  #-}
{-# LANGUAGE TypeFamilies     #-}
{-# LANGUAGE TypeOperators    #-}
{-# LANGUAGE ViewPatterns     #-}
{- |
An inlining pass.

This pass is essentially a copy of the PIR inliner, and should be KEPT IN SYNC with it.
It's hard to do this with true abstraction, so we just have to keep two copies reasonably
similar.

However, there are some differences. In the interests of making it easier to keep
things in sync, these are explicitly listed in Note [Differences from PIR inliner].
If you add another difference, please note it there! Obviously fewer differences is
better.

See Note [The problem of inlining destructors].
-}
module UntypedPlutusCore.Transform.Inline (inline, InlineHints (..)) where

import PlutusPrelude
import UntypedPlutusCore.Analysis.Usages qualified as Usages
import UntypedPlutusCore.Core.Plated
import UntypedPlutusCore.Core.Type
import UntypedPlutusCore.MkUPlc
import UntypedPlutusCore.Rename ()

import PlutusCore qualified as PLC
import PlutusCore.Builtin qualified as PLC
import PlutusCore.InlineUtils
import PlutusCore.Name
import PlutusCore.Quote

import Control.Lens hiding (Strict)
import Control.Monad.Reader
import Control.Monad.State

import Data.Map qualified as Map
import Data.Semigroup.Generic (GenericSemigroupMonoid (..))
import Witherable

{- Note [Differences from PIR inliner]

1. No types (obviously).
2. No strictness information (we only have lambda arguments, which are always strict).
3. Handling of multiple beta-reductions in one go, this is handled in PIR by a dedicated pass.
4. Don't inline lambdas with small bodies. We do this in PIR but we *probably* shouldn't really. But doing it here
is actively harmful, so we don't include it.
5. Simplistic purity analysis, in particular we don't try to be clever about builtins (should mostly be handled in PIR).
-}

newtype InlineTerm name uni fun a = Done (Term name uni fun a)

newtype TermEnv name uni fun a = TermEnv { TermEnv name uni fun a
-> UniqueMap TermUnique (InlineTerm name uni fun a)
_unTermEnv :: PLC.UniqueMap TermUnique (InlineTerm name uni fun a) }
    deriving newtype (b -> TermEnv name uni fun a -> TermEnv name uni fun a
NonEmpty (TermEnv name uni fun a) -> TermEnv name uni fun a
TermEnv name uni fun a
-> TermEnv name uni fun a -> TermEnv name uni fun a
(TermEnv name uni fun a
 -> TermEnv name uni fun a -> TermEnv name uni fun a)
-> (NonEmpty (TermEnv name uni fun a) -> TermEnv name uni fun a)
-> (forall b.
    Integral b =>
    b -> TermEnv name uni fun a -> TermEnv name uni fun a)
-> Semigroup (TermEnv name uni fun a)
forall b.
Integral b =>
b -> TermEnv name uni fun a -> TermEnv name uni fun a
forall a.
(a -> a -> a)
-> (NonEmpty a -> a)
-> (forall b. Integral b => b -> a -> a)
-> Semigroup a
forall name (uni :: * -> *) fun a.
NonEmpty (TermEnv name uni fun a) -> TermEnv name uni fun a
forall name (uni :: * -> *) fun a.
TermEnv name uni fun a
-> TermEnv name uni fun a -> TermEnv name uni fun a
forall name (uni :: * -> *) fun a b.
Integral b =>
b -> TermEnv name uni fun a -> TermEnv name uni fun a
stimes :: b -> TermEnv name uni fun a -> TermEnv name uni fun a
$cstimes :: forall name (uni :: * -> *) fun a b.
Integral b =>
b -> TermEnv name uni fun a -> TermEnv name uni fun a
sconcat :: NonEmpty (TermEnv name uni fun a) -> TermEnv name uni fun a
$csconcat :: forall name (uni :: * -> *) fun a.
NonEmpty (TermEnv name uni fun a) -> TermEnv name uni fun a
<> :: TermEnv name uni fun a
-> TermEnv name uni fun a -> TermEnv name uni fun a
$c<> :: forall name (uni :: * -> *) fun a.
TermEnv name uni fun a
-> TermEnv name uni fun a -> TermEnv name uni fun a
Semigroup, Semigroup (TermEnv name uni fun a)
TermEnv name uni fun a
Semigroup (TermEnv name uni fun a)
-> TermEnv name uni fun a
-> (TermEnv name uni fun a
    -> TermEnv name uni fun a -> TermEnv name uni fun a)
-> ([TermEnv name uni fun a] -> TermEnv name uni fun a)
-> Monoid (TermEnv name uni fun a)
[TermEnv name uni fun a] -> TermEnv name uni fun a
TermEnv name uni fun a
-> TermEnv name uni fun a -> TermEnv name uni fun a
forall a.
Semigroup a -> a -> (a -> a -> a) -> ([a] -> a) -> Monoid a
forall name (uni :: * -> *) fun a.
Semigroup (TermEnv name uni fun a)
forall name (uni :: * -> *) fun a. TermEnv name uni fun a
forall name (uni :: * -> *) fun a.
[TermEnv name uni fun a] -> TermEnv name uni fun a
forall name (uni :: * -> *) fun a.
TermEnv name uni fun a
-> TermEnv name uni fun a -> TermEnv name uni fun a
mconcat :: [TermEnv name uni fun a] -> TermEnv name uni fun a
$cmconcat :: forall name (uni :: * -> *) fun a.
[TermEnv name uni fun a] -> TermEnv name uni fun a
mappend :: TermEnv name uni fun a
-> TermEnv name uni fun a -> TermEnv name uni fun a
$cmappend :: forall name (uni :: * -> *) fun a.
TermEnv name uni fun a
-> TermEnv name uni fun a -> TermEnv name uni fun a
mempty :: TermEnv name uni fun a
$cmempty :: forall name (uni :: * -> *) fun a. TermEnv name uni fun a
$cp1Monoid :: forall name (uni :: * -> *) fun a.
Semigroup (TermEnv name uni fun a)
Monoid)

-- See Note [Differences from PIR inliner] 1
newtype Subst name uni fun a = Subst { Subst name uni fun a -> TermEnv name uni fun a
_termEnv :: TermEnv name uni fun a }
    deriving stock ((forall x. Subst name uni fun a -> Rep (Subst name uni fun a) x)
-> (forall x. Rep (Subst name uni fun a) x -> Subst name uni fun a)
-> Generic (Subst name uni fun a)
forall x. Rep (Subst name uni fun a) x -> Subst name uni fun a
forall x. Subst name uni fun a -> Rep (Subst name uni fun a) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall name (uni :: * -> *) fun a x.
Rep (Subst name uni fun a) x -> Subst name uni fun a
forall name (uni :: * -> *) fun a x.
Subst name uni fun a -> Rep (Subst name uni fun a) x
$cto :: forall name (uni :: * -> *) fun a x.
Rep (Subst name uni fun a) x -> Subst name uni fun a
$cfrom :: forall name (uni :: * -> *) fun a x.
Subst name uni fun a -> Rep (Subst name uni fun a) x
Generic)
    deriving (b -> Subst name uni fun a -> Subst name uni fun a
NonEmpty (Subst name uni fun a) -> Subst name uni fun a
Subst name uni fun a
-> Subst name uni fun a -> Subst name uni fun a
(Subst name uni fun a
 -> Subst name uni fun a -> Subst name uni fun a)
-> (NonEmpty (Subst name uni fun a) -> Subst name uni fun a)
-> (forall b.
    Integral b =>
    b -> Subst name uni fun a -> Subst name uni fun a)
-> Semigroup (Subst name uni fun a)
forall b.
Integral b =>
b -> Subst name uni fun a -> Subst name uni fun a
forall a.
(a -> a -> a)
-> (NonEmpty a -> a)
-> (forall b. Integral b => b -> a -> a)
-> Semigroup a
forall name (uni :: * -> *) fun a.
NonEmpty (Subst name uni fun a) -> Subst name uni fun a
forall name (uni :: * -> *) fun a.
Subst name uni fun a
-> Subst name uni fun a -> Subst name uni fun a
forall name (uni :: * -> *) fun a b.
Integral b =>
b -> Subst name uni fun a -> Subst name uni fun a
stimes :: b -> Subst name uni fun a -> Subst name uni fun a
$cstimes :: forall name (uni :: * -> *) fun a b.
Integral b =>
b -> Subst name uni fun a -> Subst name uni fun a
sconcat :: NonEmpty (Subst name uni fun a) -> Subst name uni fun a
$csconcat :: forall name (uni :: * -> *) fun a.
NonEmpty (Subst name uni fun a) -> Subst name uni fun a
<> :: Subst name uni fun a
-> Subst name uni fun a -> Subst name uni fun a
$c<> :: forall name (uni :: * -> *) fun a.
Subst name uni fun a
-> Subst name uni fun a -> Subst name uni fun a
Semigroup, Semigroup (Subst name uni fun a)
Subst name uni fun a
Semigroup (Subst name uni fun a)
-> Subst name uni fun a
-> (Subst name uni fun a
    -> Subst name uni fun a -> Subst name uni fun a)
-> ([Subst name uni fun a] -> Subst name uni fun a)
-> Monoid (Subst name uni fun a)
[Subst name uni fun a] -> Subst name uni fun a
Subst name uni fun a
-> Subst name uni fun a -> Subst name uni fun a
forall a.
Semigroup a -> a -> (a -> a -> a) -> ([a] -> a) -> Monoid a
forall name (uni :: * -> *) fun a. Semigroup (Subst name uni fun a)
forall name (uni :: * -> *) fun a. Subst name uni fun a
forall name (uni :: * -> *) fun a.
[Subst name uni fun a] -> Subst name uni fun a
forall name (uni :: * -> *) fun a.
Subst name uni fun a
-> Subst name uni fun a -> Subst name uni fun a
mconcat :: [Subst name uni fun a] -> Subst name uni fun a
$cmconcat :: forall name (uni :: * -> *) fun a.
[Subst name uni fun a] -> Subst name uni fun a
mappend :: Subst name uni fun a
-> Subst name uni fun a -> Subst name uni fun a
$cmappend :: forall name (uni :: * -> *) fun a.
Subst name uni fun a
-> Subst name uni fun a -> Subst name uni fun a
mempty :: Subst name uni fun a
$cmempty :: forall name (uni :: * -> *) fun a. Subst name uni fun a
$cp1Monoid :: forall name (uni :: * -> *) fun a. Semigroup (Subst name uni fun a)
Monoid) via (GenericSemigroupMonoid (Subst name uni fun a))

makeLenses ''TermEnv
makeLenses ''Subst

type ExternalConstraints name uni fun m =
    ( HasUnique name TermUnique
    , PLC.ToBuiltinMeaning uni fun
    , MonadQuote m
    )

type InliningConstraints name uni fun =
    ( HasUnique name TermUnique
    , PLC.ToBuiltinMeaning uni fun
    )

-- See Note [Differences from PIR inliner] 2
data InlineInfo name a = InlineInfo { InlineInfo name a -> Usages
_usages :: Usages.Usages, InlineInfo name a -> InlineHints name a
_hints :: InlineHints name a }

-- Using a concrete monad makes a very large difference to the performance of this module (determined from profiling)
type InlineM name uni fun a = ReaderT (InlineInfo name a) (StateT (Subst name uni fun a) Quote)

lookupTerm
    :: (HasUnique name TermUnique)
    => name
    -> Subst name uni fun a
    -> Maybe (InlineTerm name uni fun a)
lookupTerm :: name -> Subst name uni fun a -> Maybe (InlineTerm name uni fun a)
lookupTerm name
n Subst name uni fun a
subst = name
-> UniqueMap TermUnique (InlineTerm name uni fun a)
-> Maybe (InlineTerm name uni fun a)
forall name unique a.
HasUnique name unique =>
name -> UniqueMap unique a -> Maybe a
lookupName name
n (UniqueMap TermUnique (InlineTerm name uni fun a)
 -> Maybe (InlineTerm name uni fun a))
-> UniqueMap TermUnique (InlineTerm name uni fun a)
-> Maybe (InlineTerm name uni fun a)
forall a b. (a -> b) -> a -> b
$ Subst name uni fun a
subst Subst name uni fun a
-> Getting
     (UniqueMap TermUnique (InlineTerm name uni fun a))
     (Subst name uni fun a)
     (UniqueMap TermUnique (InlineTerm name uni fun a))
-> UniqueMap TermUnique (InlineTerm name uni fun a)
forall s a. s -> Getting a s a -> a
^. (TermEnv name uni fun a
 -> Const
      (UniqueMap TermUnique (InlineTerm name uni fun a))
      (TermEnv name uni fun a))
-> Subst name uni fun a
-> Const
     (UniqueMap TermUnique (InlineTerm name uni fun a))
     (Subst name uni fun a)
forall name (uni :: * -> *) fun a name (uni :: * -> *) fun a.
Iso
  (Subst name uni fun a)
  (Subst name uni fun a)
  (TermEnv name uni fun a)
  (TermEnv name uni fun a)
termEnv ((TermEnv name uni fun a
  -> Const
       (UniqueMap TermUnique (InlineTerm name uni fun a))
       (TermEnv name uni fun a))
 -> Subst name uni fun a
 -> Const
      (UniqueMap TermUnique (InlineTerm name uni fun a))
      (Subst name uni fun a))
-> ((UniqueMap TermUnique (InlineTerm name uni fun a)
     -> Const
          (UniqueMap TermUnique (InlineTerm name uni fun a))
          (UniqueMap TermUnique (InlineTerm name uni fun a)))
    -> TermEnv name uni fun a
    -> Const
         (UniqueMap TermUnique (InlineTerm name uni fun a))
         (TermEnv name uni fun a))
-> Getting
     (UniqueMap TermUnique (InlineTerm name uni fun a))
     (Subst name uni fun a)
     (UniqueMap TermUnique (InlineTerm name uni fun a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (UniqueMap TermUnique (InlineTerm name uni fun a)
 -> Const
      (UniqueMap TermUnique (InlineTerm name uni fun a))
      (UniqueMap TermUnique (InlineTerm name uni fun a)))
-> TermEnv name uni fun a
-> Const
     (UniqueMap TermUnique (InlineTerm name uni fun a))
     (TermEnv name uni fun a)
forall name (uni :: * -> *) fun a name (uni :: * -> *) fun a.
Iso
  (TermEnv name uni fun a)
  (TermEnv name uni fun a)
  (UniqueMap TermUnique (InlineTerm name uni fun a))
  (UniqueMap TermUnique (InlineTerm name uni fun a))
unTermEnv

extendTerm
    :: (HasUnique name TermUnique)
    => name
    -> InlineTerm name uni fun a
    -> Subst name uni fun a
    -> Subst name uni fun a
extendTerm :: name
-> InlineTerm name uni fun a
-> Subst name uni fun a
-> Subst name uni fun a
extendTerm name
n InlineTerm name uni fun a
clos Subst name uni fun a
subst = Subst name uni fun a
subst Subst name uni fun a
-> (Subst name uni fun a -> Subst name uni fun a)
-> Subst name uni fun a
forall a b. a -> (a -> b) -> b
& (TermEnv name uni fun a -> Identity (TermEnv name uni fun a))
-> Subst name uni fun a -> Identity (Subst name uni fun a)
forall name (uni :: * -> *) fun a name (uni :: * -> *) fun a.
Iso
  (Subst name uni fun a)
  (Subst name uni fun a)
  (TermEnv name uni fun a)
  (TermEnv name uni fun a)
termEnv ((TermEnv name uni fun a -> Identity (TermEnv name uni fun a))
 -> Subst name uni fun a -> Identity (Subst name uni fun a))
-> ((UniqueMap TermUnique (InlineTerm name uni fun a)
     -> Identity (UniqueMap TermUnique (InlineTerm name uni fun a)))
    -> TermEnv name uni fun a -> Identity (TermEnv name uni fun a))
-> (UniqueMap TermUnique (InlineTerm name uni fun a)
    -> Identity (UniqueMap TermUnique (InlineTerm name uni fun a)))
-> Subst name uni fun a
-> Identity (Subst name uni fun a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (UniqueMap TermUnique (InlineTerm name uni fun a)
 -> Identity (UniqueMap TermUnique (InlineTerm name uni fun a)))
-> TermEnv name uni fun a -> Identity (TermEnv name uni fun a)
forall name (uni :: * -> *) fun a name (uni :: * -> *) fun a.
Iso
  (TermEnv name uni fun a)
  (TermEnv name uni fun a)
  (UniqueMap TermUnique (InlineTerm name uni fun a))
  (UniqueMap TermUnique (InlineTerm name uni fun a))
unTermEnv ((UniqueMap TermUnique (InlineTerm name uni fun a)
  -> Identity (UniqueMap TermUnique (InlineTerm name uni fun a)))
 -> Subst name uni fun a -> Identity (Subst name uni fun a))
-> (UniqueMap TermUnique (InlineTerm name uni fun a)
    -> UniqueMap TermUnique (InlineTerm name uni fun a))
-> Subst name uni fun a
-> Subst name uni fun a
forall s t a b. ASetter s t a b -> (a -> b) -> s -> t
%~ name
-> InlineTerm name uni fun a
-> UniqueMap TermUnique (InlineTerm name uni fun a)
-> UniqueMap TermUnique (InlineTerm name uni fun a)
forall name unique a.
HasUnique name unique =>
name -> a -> UniqueMap unique a -> UniqueMap unique a
insertByName name
n InlineTerm name uni fun a
clos

-- | Inline simple bindings. Relies on global uniqueness, and preserves it.
-- See Note [Inlining and global uniqueness]
inline
    :: forall name uni fun m a
    . ExternalConstraints name uni fun m
    => InlineHints name a
    -> Term name uni fun a
    -> m (Term name uni fun a)
inline :: InlineHints name a
-> Term name uni fun a -> m (Term name uni fun a)
inline InlineHints name a
hints Term name uni fun a
t = let
        inlineInfo :: InlineInfo name a
        inlineInfo :: InlineInfo name a
inlineInfo = Usages -> InlineHints name a -> InlineInfo name a
forall name a. Usages -> InlineHints name a -> InlineInfo name a
InlineInfo Usages
usgs InlineHints name a
hints
        usgs :: Map.Map Unique Int
        usgs :: Usages
usgs = Term name uni fun a -> Usages
forall name (uni :: * -> *) fun a.
HasUnique name TermUnique =>
Term name uni fun a -> Usages
Usages.runTermUsages Term name uni fun a
t
    in Quote (Term name uni fun a) -> m (Term name uni fun a)
forall (m :: * -> *) a. MonadQuote m => Quote a -> m a
liftQuote (Quote (Term name uni fun a) -> m (Term name uni fun a))
-> Quote (Term name uni fun a) -> m (Term name uni fun a)
forall a b. (a -> b) -> a -> b
$ (StateT (Subst name uni fun a) Quote (Term name uni fun a)
 -> Subst name uni fun a -> Quote (Term name uni fun a))
-> Subst name uni fun a
-> StateT (Subst name uni fun a) Quote (Term name uni fun a)
-> Quote (Term name uni fun a)
forall a b c. (a -> b -> c) -> b -> a -> c
flip StateT (Subst name uni fun a) Quote (Term name uni fun a)
-> Subst name uni fun a -> Quote (Term name uni fun a)
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT Subst name uni fun a
forall a. Monoid a => a
mempty (StateT (Subst name uni fun a) Quote (Term name uni fun a)
 -> Quote (Term name uni fun a))
-> StateT (Subst name uni fun a) Quote (Term name uni fun a)
-> Quote (Term name uni fun a)
forall a b. (a -> b) -> a -> b
$ (ReaderT
   (InlineInfo name a)
   (StateT (Subst name uni fun a) Quote)
   (Term name uni fun a)
 -> InlineInfo name a
 -> StateT (Subst name uni fun a) Quote (Term name uni fun a))
-> InlineInfo name a
-> ReaderT
     (InlineInfo name a)
     (StateT (Subst name uni fun a) Quote)
     (Term name uni fun a)
-> StateT (Subst name uni fun a) Quote (Term name uni fun a)
forall a b c. (a -> b -> c) -> b -> a -> c
flip ReaderT
  (InlineInfo name a)
  (StateT (Subst name uni fun a) Quote)
  (Term name uni fun a)
-> InlineInfo name a
-> StateT (Subst name uni fun a) Quote (Term name uni fun a)
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT InlineInfo name a
inlineInfo (ReaderT
   (InlineInfo name a)
   (StateT (Subst name uni fun a) Quote)
   (Term name uni fun a)
 -> StateT (Subst name uni fun a) Quote (Term name uni fun a))
-> ReaderT
     (InlineInfo name a)
     (StateT (Subst name uni fun a) Quote)
     (Term name uni fun a)
-> StateT (Subst name uni fun a) Quote (Term name uni fun a)
forall a b. (a -> b) -> a -> b
$ Term name uni fun a
-> ReaderT
     (InlineInfo name a)
     (StateT (Subst name uni fun a) Quote)
     (Term name uni fun a)
forall name (uni :: * -> *) fun a.
InliningConstraints name uni fun =>
Term name uni fun a -> InlineM name uni fun a (Term name uni fun a)
processTerm Term name uni fun a
t

-- See Note [Differences from PIR inliner] 3
{-| Extract the list of applications from a term, a bit like a "multi-beta" reduction.

Some examples will help:
[(\x . t) a] -> Just ([(x, a)], t)

[[[(\x . (\y . (\z . t))) a] b] c] -> Just ([(x, a), (y, b), (z, c)]) t)

[[(\x . t) a] b] -> Nothing
-}
extractApps :: Term name uni fun a -> Maybe ([UTermDef name uni fun a], Term name uni fun a)
extractApps :: Term name uni fun a
-> Maybe ([UTermDef name uni fun a], Term name uni fun a)
extractApps = [Term name uni fun a]
-> Term name uni fun a
-> Maybe ([UTermDef name uni fun a], Term name uni fun a)
forall name (uni :: * -> *) fun ann.
[Term name uni fun ann]
-> Term name uni fun ann
-> Maybe
     ([Def (UVarDecl name ann) (Term name uni fun ann)],
      Term name uni fun ann)
collectArgs []
  where
      collectArgs :: [Term name uni fun ann]
-> Term name uni fun ann
-> Maybe
     ([Def (UVarDecl name ann) (Term name uni fun ann)],
      Term name uni fun ann)
collectArgs [Term name uni fun ann]
argStack (Apply ann
_ Term name uni fun ann
f Term name uni fun ann
arg) = [Term name uni fun ann]
-> Term name uni fun ann
-> Maybe
     ([Def (UVarDecl name ann) (Term name uni fun ann)],
      Term name uni fun ann)
collectArgs (Term name uni fun ann
argTerm name uni fun ann
-> [Term name uni fun ann] -> [Term name uni fun ann]
forall a. a -> [a] -> [a]
:[Term name uni fun ann]
argStack) Term name uni fun ann
f
      collectArgs [Term name uni fun ann]
argStack Term name uni fun ann
t               = [Term name uni fun ann]
-> [Def (UVarDecl name ann) (Term name uni fun ann)]
-> Term name uni fun ann
-> Maybe
     ([Def (UVarDecl name ann) (Term name uni fun ann)],
      Term name uni fun ann)
forall val name ann (uni :: * -> *) fun.
[val]
-> [Def (UVarDecl name ann) val]
-> Term name uni fun ann
-> Maybe ([Def (UVarDecl name ann) val], Term name uni fun ann)
matchArgs [Term name uni fun ann]
argStack [] Term name uni fun ann
t
      matchArgs :: [val]
-> [Def (UVarDecl name ann) val]
-> Term name uni fun ann
-> Maybe ([Def (UVarDecl name ann) val], Term name uni fun ann)
matchArgs (val
arg:[val]
rest) [Def (UVarDecl name ann) val]
acc (LamAbs ann
a name
n Term name uni fun ann
body) = [val]
-> [Def (UVarDecl name ann) val]
-> Term name uni fun ann
-> Maybe ([Def (UVarDecl name ann) val], Term name uni fun ann)
matchArgs [val]
rest (UVarDecl name ann -> val -> Def (UVarDecl name ann) val
forall var val. var -> val -> Def var val
Def (ann -> name -> UVarDecl name ann
forall name ann. ann -> name -> UVarDecl name ann
UVarDecl ann
a name
n) val
argDef (UVarDecl name ann) val
-> [Def (UVarDecl name ann) val] -> [Def (UVarDecl name ann) val]
forall a. a -> [a] -> [a]
:[Def (UVarDecl name ann) val]
acc) Term name uni fun ann
body
      matchArgs []         [Def (UVarDecl name ann) val]
acc Term name uni fun ann
t                 = if [Def (UVarDecl name ann) val] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Def (UVarDecl name ann) val]
acc then Maybe ([Def (UVarDecl name ann) val], Term name uni fun ann)
forall a. Maybe a
Nothing else ([Def (UVarDecl name ann) val], Term name uni fun ann)
-> Maybe ([Def (UVarDecl name ann) val], Term name uni fun ann)
forall a. a -> Maybe a
Just ([Def (UVarDecl name ann) val] -> [Def (UVarDecl name ann) val]
forall a. [a] -> [a]
reverse [Def (UVarDecl name ann) val]
acc, Term name uni fun ann
t)
      matchArgs (val
_:[val]
_)      [Def (UVarDecl name ann) val]
_   Term name uni fun ann
_                 = Maybe ([Def (UVarDecl name ann) val], Term name uni fun ann)
forall a. Maybe a
Nothing

-- | The inverse of 'extractApps'.
restoreApps :: [UTermDef name uni fun a] -> Term name uni fun a -> Term name uni fun a
restoreApps :: [UTermDef name uni fun a]
-> Term name uni fun a -> Term name uni fun a
restoreApps [UTermDef name uni fun a]
defs Term name uni fun a
t = [Term name uni fun a]
-> Term name uni fun a
-> [UTermDef name uni fun a]
-> Term name uni fun a
forall name (uni :: * -> *) fun ann.
[Term name uni fun ann]
-> Term name uni fun ann
-> [Def (UVarDecl name ann) (Term name uni fun ann)]
-> Term name uni fun ann
makeLams [] Term name uni fun a
t ([UTermDef name uni fun a] -> [UTermDef name uni fun a]
forall a. [a] -> [a]
reverse [UTermDef name uni fun a]
defs)
  where
      makeLams :: [Term name uni fun ann]
-> Term name uni fun ann
-> [Def (UVarDecl name ann) (Term name uni fun ann)]
-> Term name uni fun ann
makeLams [Term name uni fun ann]
args Term name uni fun ann
acc (Def (UVarDecl ann
a name
n) Term name uni fun ann
rhs:[Def (UVarDecl name ann) (Term name uni fun ann)]
rest) = [Term name uni fun ann]
-> Term name uni fun ann
-> [Def (UVarDecl name ann) (Term name uni fun ann)]
-> Term name uni fun ann
makeLams (Term name uni fun ann
rhsTerm name uni fun ann
-> [Term name uni fun ann] -> [Term name uni fun ann]
forall a. a -> [a] -> [a]
:[Term name uni fun ann]
args) (ann -> name -> Term name uni fun ann -> Term name uni fun ann
forall name (uni :: * -> *) fun ann.
ann -> name -> Term name uni fun ann -> Term name uni fun ann
LamAbs ann
a name
n Term name uni fun ann
acc) [Def (UVarDecl name ann) (Term name uni fun ann)]
rest
      makeLams [Term name uni fun ann]
args Term name uni fun ann
acc []                            = [Term name uni fun ann]
-> Term name uni fun ann -> Term name uni fun ann
forall name (uni :: * -> *) fun ann.
[Term name uni fun ann]
-> Term name uni fun ann -> Term name uni fun ann
makeApps [Term name uni fun ann]
args Term name uni fun ann
acc
      -- This isn't the best annotation, but it will do
      makeApps :: [Term name uni fun ann]
-> Term name uni fun ann -> Term name uni fun ann
makeApps (Term name uni fun ann
arg:[Term name uni fun ann]
args) Term name uni fun ann
acc = [Term name uni fun ann]
-> Term name uni fun ann -> Term name uni fun ann
makeApps [Term name uni fun ann]
args (ann
-> Term name uni fun ann
-> Term name uni fun ann
-> Term name uni fun ann
forall name (uni :: * -> *) fun ann.
ann
-> Term name uni fun ann
-> Term name uni fun ann
-> Term name uni fun ann
Apply (Term name uni fun ann -> ann
forall name (uni :: * -> *) fun ann. Term name uni fun ann -> ann
termAnn Term name uni fun ann
acc) Term name uni fun ann
acc Term name uni fun ann
arg)
      makeApps [] Term name uni fun ann
acc         = Term name uni fun ann
acc

processTerm
    :: forall name uni fun a. InliningConstraints name uni fun
    => Term name uni fun a
    -> InlineM name uni fun a (Term name uni fun a)
processTerm :: Term name uni fun a -> InlineM name uni fun a (Term name uni fun a)
processTerm = Term name uni fun a -> InlineM name uni fun a (Term name uni fun a)
handleTerm where
    handleTerm :: Term name uni fun a -> InlineM name uni fun a (Term name uni fun a)
    handleTerm :: Term name uni fun a -> InlineM name uni fun a (Term name uni fun a)
handleTerm = \case
        v :: Term name uni fun a
v@(Var a
_ name
n) -> Term name uni fun a
-> Maybe (Term name uni fun a) -> Term name uni fun a
forall a. a -> Maybe a -> a
fromMaybe Term name uni fun a
v (Maybe (Term name uni fun a) -> Term name uni fun a)
-> ReaderT
     (InlineInfo name a)
     (StateT (Subst name uni fun a) Quote)
     (Maybe (Term name uni fun a))
-> InlineM name uni fun a (Term name uni fun a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> name
-> ReaderT
     (InlineInfo name a)
     (StateT (Subst name uni fun a) Quote)
     (Maybe (Term name uni fun a))
substName name
n
        -- See Note [Differences from PIR inliner] 3
        (Term name uni fun a
-> Maybe ([UTermDef name uni fun a], Term name uni fun a)
forall name (uni :: * -> *) fun a.
Term name uni fun a
-> Maybe ([UTermDef name uni fun a], Term name uni fun a)
extractApps -> Just ([UTermDef name uni fun a]
bs, Term name uni fun a
t)) -> do
            [UTermDef name uni fun a]
bs' <- (UTermDef name uni fun a
 -> ReaderT
      (InlineInfo name a)
      (StateT (Subst name uni fun a) Quote)
      (Maybe (UTermDef name uni fun a)))
-> [UTermDef name uni fun a]
-> ReaderT
     (InlineInfo name a)
     (StateT (Subst name uni fun a) Quote)
     [UTermDef name uni fun a]
forall (t :: * -> *) (f :: * -> *) a b.
(Witherable t, Applicative f) =>
(a -> f (Maybe b)) -> t a -> f (t b)
wither UTermDef name uni fun a
-> ReaderT
     (InlineInfo name a)
     (StateT (Subst name uni fun a) Quote)
     (Maybe (UTermDef name uni fun a))
forall name (uni :: * -> *) fun a.
InliningConstraints name uni fun =>
UTermDef name uni fun a
-> InlineM name uni fun a (Maybe (UTermDef name uni fun a))
processSingleBinding [UTermDef name uni fun a]
bs
            Term name uni fun a
t' <- Term name uni fun a -> InlineM name uni fun a (Term name uni fun a)
forall name (uni :: * -> *) fun a.
InliningConstraints name uni fun =>
Term name uni fun a -> InlineM name uni fun a (Term name uni fun a)
processTerm Term name uni fun a
t
            Term name uni fun a -> InlineM name uni fun a (Term name uni fun a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Term name uni fun a
 -> InlineM name uni fun a (Term name uni fun a))
-> Term name uni fun a
-> InlineM name uni fun a (Term name uni fun a)
forall a b. (a -> b) -> a -> b
$ [UTermDef name uni fun a]
-> Term name uni fun a -> Term name uni fun a
forall name (uni :: * -> *) fun a.
[UTermDef name uni fun a]
-> Term name uni fun a -> Term name uni fun a
restoreApps [UTermDef name uni fun a]
bs' Term name uni fun a
t'
        Term name uni fun a
t -> LensLike
  (WrappedMonad
     (ReaderT
        (InlineInfo name a) (StateT (Subst name uni fun a) Quote)))
  (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
    -> InlineM name uni fun a (Term name uni fun a))
-> InlineM name uni fun a (Term name uni fun a)
forall (m :: * -> *) s t a b.
LensLike (WrappedMonad m) s t a b -> s -> (a -> m b) -> m t
forMOf LensLike
  (WrappedMonad
     (ReaderT
        (InlineInfo name a) (StateT (Subst name uni fun a) Quote)))
  (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.
Traversal' (Term name uni fun ann) (Term name uni fun ann)
termSubterms Term name uni fun a
t Term name uni fun a -> InlineM name uni fun a (Term name uni fun a)
forall name (uni :: * -> *) fun a.
InliningConstraints name uni fun =>
Term name uni fun a -> InlineM name uni fun a (Term name uni fun a)
processTerm

    -- See Note [Renaming strategy]
    substName :: name -> InlineM name uni fun a (Maybe (Term name uni fun a))
    substName :: name
-> ReaderT
     (InlineInfo name a)
     (StateT (Subst name uni fun a) Quote)
     (Maybe (Term name uni fun a))
substName name
name = (Subst name uni fun a -> Maybe (InlineTerm name uni fun a))
-> ReaderT
     (InlineInfo name a)
     (StateT (Subst name uni fun a) Quote)
     (Maybe (InlineTerm name uni fun a))
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (name -> Subst name uni fun a -> Maybe (InlineTerm name uni fun a)
forall name (uni :: * -> *) fun a.
HasUnique name TermUnique =>
name -> Subst name uni fun a -> Maybe (InlineTerm name uni fun a)
lookupTerm name
name) ReaderT
  (InlineInfo name a)
  (StateT (Subst name uni fun a) Quote)
  (Maybe (InlineTerm name uni fun a))
-> (Maybe (InlineTerm name uni fun a)
    -> ReaderT
         (InlineInfo name a)
         (StateT (Subst name uni fun a) Quote)
         (Maybe (Term name uni fun a)))
-> ReaderT
     (InlineInfo name a)
     (StateT (Subst name uni fun a) Quote)
     (Maybe (Term name uni fun a))
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (InlineTerm name uni fun a
 -> InlineM name uni fun a (Term name uni fun a))
-> Maybe (InlineTerm name uni fun a)
-> ReaderT
     (InlineInfo name a)
     (StateT (Subst name uni fun a) Quote)
     (Maybe (Term name uni fun a))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse InlineTerm name uni fun a
-> InlineM name uni fun a (Term name uni fun a)
renameTerm
    -- See Note [Inlining approach and 'Secrets of the GHC Inliner']
    renameTerm :: InlineTerm name uni fun a -> InlineM name uni fun a (Term name uni fun a)
    renameTerm :: InlineTerm name uni fun a
-> InlineM name uni fun a (Term name uni fun a)
renameTerm = \case
        -- Already processed term, just rename and put it in, don't do any
        -- further optimization here.
        Done Term name uni fun a
t -> Term name uni fun a -> InlineM name uni fun a (Term name uni fun a)
forall a (m :: * -> *). (Rename a, MonadQuote m) => a -> m a
PLC.rename Term name uni fun a
t

processSingleBinding :: InliningConstraints name uni fun => UTermDef name uni fun a -> InlineM name uni fun a (Maybe (UTermDef name uni fun a))
processSingleBinding :: UTermDef name uni fun a
-> InlineM name uni fun a (Maybe (UTermDef name uni fun a))
processSingleBinding (Def vd :: UVarDecl name a
vd@(UVarDecl a
a name
n) Term name uni fun a
rhs) = do
    Maybe (Term name uni fun a)
rhs' <- a
-> name
-> Term name uni fun a
-> InlineM name uni fun a (Maybe (Term name uni fun a))
forall name (uni :: * -> *) fun a.
InliningConstraints name uni fun =>
a
-> name
-> Term name uni fun a
-> InlineM name uni fun a (Maybe (Term name uni fun a))
maybeAddSubst a
a name
n Term name uni fun a
rhs
    Maybe (UTermDef name uni fun a)
-> InlineM name uni fun a (Maybe (UTermDef name uni fun a))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (UTermDef name uni fun a)
 -> InlineM name uni fun a (Maybe (UTermDef name uni fun a)))
-> Maybe (UTermDef name uni fun a)
-> InlineM name uni fun a (Maybe (UTermDef name uni fun a))
forall a b. (a -> b) -> a -> b
$ UVarDecl name a -> Term name uni fun a -> UTermDef name uni fun a
forall var val. var -> val -> Def var val
Def UVarDecl name a
vd (Term name uni fun a -> UTermDef name uni fun a)
-> Maybe (Term name uni fun a) -> Maybe (UTermDef name uni fun a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (Term name uni fun a)
rhs'

-- NOTE:  Nothing means that we are inlining the term:
--   * we have extended the substitution, and
--   * we are removing the binding (hence we return Nothing)
maybeAddSubst
    :: forall name uni fun a. InliningConstraints name uni fun
    => a
    -> name
    -> Term name uni fun a
    -> InlineM name uni fun a (Maybe (Term name uni fun a))
maybeAddSubst :: a
-> name
-> Term name uni fun a
-> InlineM name uni fun a (Maybe (Term name uni fun a))
maybeAddSubst a
a name
n Term name uni fun a
rhs = do
    Term name uni fun a
rhs' <- Term name uni fun a -> InlineM name uni fun a (Term name uni fun a)
forall name (uni :: * -> *) fun a.
InliningConstraints name uni fun =>
Term name uni fun a -> InlineM name uni fun a (Term name uni fun a)
processTerm Term name uni fun a
rhs

    -- Check whether we've been told specifically to inline this
    InlineHints name a
hints <- (InlineInfo name a -> InlineHints name a)
-> ReaderT
     (InlineInfo name a)
     (StateT (Subst name uni fun a) Quote)
     (InlineHints name a)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks InlineInfo name a -> InlineHints name a
forall name a. InlineInfo name a -> InlineHints name a
_hints
    let hinted :: Bool
hinted = InlineHints name a -> a -> name -> Bool
forall name a. InlineHints name a -> a -> name -> Bool
shouldInline InlineHints name a
hints a
a name
n

    Bool
preUnconditional <- Term name uni fun a -> InlineM name uni fun a Bool
preInlineUnconditional Term name uni fun a
rhs'
    if Bool
preUnconditional
    then InlineTerm name uni fun a
-> InlineM name uni fun a (Maybe (Term name uni fun a))
forall b.
InlineTerm name uni fun a -> InlineM name uni fun a (Maybe b)
extendAndDrop (Term name uni fun a -> InlineTerm name uni fun a
forall name (uni :: * -> *) fun a.
Term name uni fun a -> InlineTerm name uni fun a
Done Term name uni fun a
rhs')
    else do
        -- See Note [Inlining approach and 'Secrets of the GHC Inliner']
        Bool
postUnconditional <- Term name uni fun a -> InlineM name uni fun a Bool
postInlineUnconditional Term name uni fun a
rhs'
        if Bool
hinted Bool -> Bool -> Bool
|| Bool
postUnconditional
        then InlineTerm name uni fun a
-> InlineM name uni fun a (Maybe (Term name uni fun a))
forall b.
InlineTerm name uni fun a -> InlineM name uni fun a (Maybe b)
extendAndDrop (Term name uni fun a -> InlineTerm name uni fun a
forall name (uni :: * -> *) fun a.
Term name uni fun a -> InlineTerm name uni fun a
Done Term name uni fun a
rhs')
        else Maybe (Term name uni fun a)
-> InlineM name uni fun a (Maybe (Term name uni fun a))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (Term name uni fun a)
 -> InlineM name uni fun a (Maybe (Term name uni fun a)))
-> Maybe (Term name uni fun a)
-> InlineM name uni fun a (Maybe (Term name uni fun a))
forall a b. (a -> b) -> a -> b
$ Term name uni fun a -> Maybe (Term name uni fun a)
forall a. a -> Maybe a
Just Term name uni fun a
rhs'
    where
        extendAndDrop :: forall b . InlineTerm name uni fun a -> InlineM name uni fun a (Maybe b)
        extendAndDrop :: InlineTerm name uni fun a -> InlineM name uni fun a (Maybe b)
extendAndDrop InlineTerm name uni fun a
t = (Subst name uni fun a -> Subst name uni fun a)
-> ReaderT
     (InlineInfo name a) (StateT (Subst name uni fun a) Quote) ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
modify' (name
-> InlineTerm name uni fun a
-> Subst name uni fun a
-> Subst name uni fun a
forall name (uni :: * -> *) fun a.
HasUnique name TermUnique =>
name
-> InlineTerm name uni fun a
-> Subst name uni fun a
-> Subst name uni fun a
extendTerm name
n InlineTerm name uni fun a
t) ReaderT
  (InlineInfo name a) (StateT (Subst name uni fun a) Quote) ()
-> InlineM name uni fun a (Maybe b)
-> InlineM name uni fun a (Maybe b)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Maybe b -> InlineM name uni fun a (Maybe b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe b
forall a. Maybe a
Nothing

        checkPurity :: Term name uni fun a -> InlineM name uni fun a Bool
        checkPurity :: Term name uni fun a -> InlineM name uni fun a Bool
checkPurity Term name uni fun a
t = Bool -> InlineM name uni fun a Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> InlineM name uni fun a Bool)
-> Bool -> InlineM name uni fun a Bool
forall a b. (a -> b) -> a -> b
$ Term name uni fun a -> Bool
forall name (uni :: * -> *) fun a. Term name uni fun a -> Bool
isPure Term name uni fun a
t

        preInlineUnconditional :: Term name uni fun a -> InlineM name uni fun a Bool
        preInlineUnconditional :: Term name uni fun a -> InlineM name uni fun a Bool
preInlineUnconditional Term name uni fun a
t = do
            Usages
usgs <- (InlineInfo name a -> Usages)
-> ReaderT
     (InlineInfo name a) (StateT (Subst name uni fun a) Quote) Usages
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks InlineInfo name a -> Usages
forall name a. InlineInfo name a -> Usages
_usages
            -- 'inlining' terms used 0 times is a cheap way to remove dead code while we're here
            let termUsedAtMostOnce :: Bool
termUsedAtMostOnce = name -> Usages -> Int
forall n unique. HasUnique n unique => n -> Usages -> Int
Usages.getUsageCount name
n Usages
usgs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1
            -- See Note [Inlining and purity]
            Bool
termIsPure <- Term name uni fun a -> InlineM name uni fun a Bool
checkPurity Term name uni fun a
t
            Bool -> InlineM name uni fun a Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> InlineM name uni fun a Bool)
-> Bool -> InlineM name uni fun a Bool
forall a b. (a -> b) -> a -> b
$ Bool
termUsedAtMostOnce Bool -> Bool -> Bool
&& Bool
termIsPure

        -- | Should we inline? Should only inline things that won't duplicate work or code.
        -- See Note [Inlining approach and 'Secrets of the GHC Inliner']
        postInlineUnconditional ::  Term name uni fun a -> InlineM name uni fun a Bool
        postInlineUnconditional :: Term name uni fun a -> InlineM name uni fun a Bool
postInlineUnconditional Term name uni fun a
t = do
            -- See Note [Inlining criteria]
            let acceptable :: Bool
acceptable = Term name uni fun a -> Bool
forall name (uni :: * -> *) fun a. Term name uni fun a -> Bool
costIsAcceptable Term name uni fun a
t Bool -> Bool -> Bool
&& Term name uni fun a -> Bool
forall name (uni :: * -> *) fun a. Term name uni fun a -> Bool
sizeIsAcceptable Term name uni fun a
t
            -- See Note [Inlining and purity]
            Bool
termIsPure <- Term name uni fun a -> InlineM name uni fun a Bool
checkPurity Term name uni fun a
t
            Bool -> InlineM name uni fun a Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> InlineM name uni fun a Bool)
-> Bool -> InlineM name uni fun a Bool
forall a b. (a -> b) -> a -> b
$ Bool
acceptable Bool -> Bool -> Bool
&& Bool
termIsPure

-- | Is the cost increase (in terms of evaluation work) of inlining a variable whose RHS is
-- the given term acceptable?
costIsAcceptable :: Term name uni fun a -> Bool
costIsAcceptable :: Term name uni fun a -> Bool
costIsAcceptable = \case
  Builtin{}  -> Bool
True
  Var{}      -> Bool
True
  Constant{} -> Bool
True
  Error{}    -> Bool
True
  -- This will mean that we create closures at each use site instead of
  -- once, but that's a very low cost which we're okay rounding to 0.
  LamAbs{}   -> Bool
True

  Apply{}    -> Bool
False

  Force{}    -> Bool
True
  Delay{}    -> Bool
True

-- | Is the size increase (in the AST) of inlining a variable whose RHS is
-- the given term acceptable?
sizeIsAcceptable :: Term name uni fun a -> Bool
sizeIsAcceptable :: Term name uni fun a -> Bool
sizeIsAcceptable = \case
  Builtin{}  -> Bool
True
  Var{}      -> Bool
True
  Error{}    -> Bool
True

  -- See Note [Differences from PIR inliner] 4
  LamAbs{}   -> Bool
False

  -- Constants can be big! We could check the size here and inline if they're
  -- small, but probably not worth it
  Constant{} -> Bool
False
  Apply{}    -> Bool
False

  Force a
_ Term name uni fun a
t  -> Term name uni fun a -> Bool
forall name (uni :: * -> *) fun a. Term name uni fun a -> Bool
sizeIsAcceptable Term name uni fun a
t
  Delay a
_ Term name uni fun a
t  -> Term name uni fun a -> Bool
forall name (uni :: * -> *) fun a. Term name uni fun a -> Bool
sizeIsAcceptable Term name uni fun a
t

isPure :: Term name uni fun a -> Bool
isPure :: Term name uni fun a -> Bool
isPure = Term name uni fun a -> Bool
forall name (uni :: * -> *) fun a. Term name uni fun a -> Bool
go
    where
        go :: Term name uni fun ann -> Bool
go = \case
            Var {}      -> Bool
True
            -- These are syntactically values that won't reduce further
            LamAbs {}   -> Bool
True
            Constant {} -> Bool
True
            Delay {}    -> Bool
True

            -- See Note [Differences from PIR inliner] 5
            Term name uni fun ann
_           -> Bool
False