{-# LANGUAGE Safe #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
-- | The INLINE-attribute inliner (the retired Crust pass's successor):
--   every occurrence of an INLINE-annotated definition is replaced by its
--   body (as a lambda telescope over its parameters — application sites
--   become beta redexes for the downstream partial evaluator), with every
--   inserted copy refreshed through the audited clone primitive. Runs on
--   monomorphic programs (after 'ReWire.Eidos.Spec.specialize'): inlining
--   under a type-argument spine would strand the arguments on a
--   non-variable head.
--
--   INLINE definitions referencing other INLINE definitions are expanded
--   to closed form first (depth-first, memoized); a reference cycle among
--   them is rejected with the retired pass's diagnostic. The definitions
--   themselves are kept (dead ones fall to the downstream purge).
module ReWire.Eidos.Inline (inlineAnnotated) where

import ReWire.Error (AstError, MonadError, failAt)
import ReWire.Eidos.Subst (nextUniq, occCounts, substVarsRefreshing)
import ReWire.Eidos.Syntax

import Control.Monad (foldM)
import Control.Monad.State.Strict (StateT, evalStateT)

import qualified Data.IntMap.Strict as IM
import qualified Data.IntSet        as IS

inlineAnnotated :: forall m. MonadError AstError m => Program -> m Program
inlineAnnotated :: forall (m :: * -> *). MonadError AstError m => Program -> m Program
inlineAnnotated p :: Program
p@(Program [DataDefn]
datas [Defn]
defns Id
top) = StateT Uniq m Program -> Uniq -> m Program
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT StateT Uniq m Program
go (Uniq -> m Program) -> Uniq -> m Program
forall a b. (a -> b) -> a -> b
$ Program -> Uniq
forall a. Data a => a -> Uniq
nextUniq Program
p
      where inls :: IM.IntMap Defn
            inls :: IntMap Defn
inls = [(Uniq, Defn)] -> IntMap Defn
forall a. [(Uniq, a)] -> IntMap a
IM.fromList [ (Id -> Uniq
idUniq (Id -> Uniq) -> Id -> Uniq
forall a b. (a -> b) -> a -> b
$ Defn -> Id
defnId Defn
d, Defn
d) | Defn
d <- [Defn]
defns, Defn -> Maybe DefnAttr
defnAttr Defn
d Maybe DefnAttr -> Maybe DefnAttr -> Bool
forall a. Eq a => a -> a -> Bool
== DefnAttr -> Maybe DefnAttr
forall a. a -> Maybe a
Just DefnAttr
Inline ]

            go :: StateT Uniq m Program
            go :: StateT Uniq m Program
go = do
                  table  <- (IntMap Exp -> Uniq -> StateT Uniq m (IntMap Exp))
-> IntMap Exp -> [Uniq] -> StateT Uniq m (IntMap Exp)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (\ IntMap Exp
acc Uniq
u -> IntMap Exp -> IntSet -> Uniq -> StateT Uniq m (IntMap Exp)
expand IntMap Exp
acc IntSet
forall a. Monoid a => a
mempty Uniq
u) IntMap Exp
forall a. Monoid a => a
mempty ([Uniq] -> StateT Uniq m (IntMap Exp))
-> [Uniq] -> StateT Uniq m (IntMap Exp)
forall a b. (a -> b) -> a -> b
$ IntMap Defn -> [Uniq]
forall a. IntMap a -> [Uniq]
IM.keys IntMap Defn
inls
                  defns' <- mapM (\ Defn
d -> (\ Exp
b -> Defn
d { defnBody = b }) (Exp -> Defn) -> StateT Uniq m Exp -> StateT Uniq m Defn
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IntMap Exp -> Exp -> StateT Uniq m Exp
forall (m :: * -> *).
MonadState Uniq m =>
IntMap Exp -> Exp -> m Exp
substVarsRefreshing IntMap Exp
table (Defn -> Exp
defnBody Defn
d)) defns
                  pure $ Program datas defns' top

            -- The definition as a substitution payload: a lambda telescope
            -- over its parameters.
            payload :: Defn -> Exp
            payload :: Defn -> Exp
payload Defn
d = (Id -> Exp -> Exp) -> Exp -> [Id] -> Exp
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (Annote -> Id -> Exp -> Exp
Lam (Annote -> Id -> Exp -> Exp) -> Annote -> Id -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ Defn -> Annote
defnAnnote Defn
d) (Defn -> Exp
defnBody Defn
d) ([Id] -> Exp) -> [Id] -> Exp
forall a b. (a -> b) -> a -> b
$ Defn -> [Id]
defnParams Defn
d

            -- Expand one INLINE definition to closed form (no INLINE
            -- references remain), depth-first over its INLINE dependencies.
            expand :: IM.IntMap Exp -> IS.IntSet -> Uniq -> StateT Uniq m (IM.IntMap Exp)
            expand :: IntMap Exp -> IntSet -> Uniq -> StateT Uniq m (IntMap Exp)
expand IntMap Exp
table IntSet
stack Uniq
u
                  | Uniq -> IntMap Exp -> Bool
forall a. Uniq -> IntMap a -> Bool
IM.member Uniq
u IntMap Exp
table = IntMap Exp -> StateT Uniq m (IntMap Exp)
forall a. a -> StateT Uniq m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure IntMap Exp
table
                  | Uniq -> IntSet -> Bool
IS.member Uniq
u IntSet
stack = Annote -> Text -> StateT Uniq m (IntMap Exp)
forall (m :: * -> *) an a.
(MonadError AstError m, Annotation an) =>
an -> Text -> m a
failAt (Defn -> Annote
defnAnnote Defn
d) Text
"INLINE definition expansion not terminating (mutually recursive definitions?)."
                  | Bool
otherwise         = do
                        let deps :: [Uniq]
deps = IntSet -> [Uniq]
IS.toList (IntSet -> [Uniq]) -> IntSet -> [Uniq]
forall a b. (a -> b) -> a -> b
$ IntMap Uniq -> IntSet
forall a. IntMap a -> IntSet
IM.keysSet (Exp -> IntMap Uniq
occCounts (Exp -> IntMap Uniq) -> Exp -> IntMap Uniq
forall a b. (a -> b) -> a -> b
$ Defn -> Exp
payload Defn
d) IntSet -> IntSet -> IntSet
`IS.intersection` IntMap Defn -> IntSet
forall a. IntMap a -> IntSet
IM.keysSet IntMap Defn
inls
                        table' <- (IntMap Exp -> Uniq -> StateT Uniq m (IntMap Exp))
-> IntMap Exp -> [Uniq] -> StateT Uniq m (IntMap Exp)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM (\ IntMap Exp
acc Uniq
v -> IntMap Exp -> IntSet -> Uniq -> StateT Uniq m (IntMap Exp)
expand IntMap Exp
acc (Uniq -> IntSet -> IntSet
IS.insert Uniq
u IntSet
stack) Uniq
v) IntMap Exp
table [Uniq]
deps
                        e'     <- substVarsRefreshing (IM.restrictKeys table' $ IS.fromList deps) $ payload d
                        pure $ IM.insert u e' table'
                  where d :: Defn
d = IntMap Defn
inls IntMap Defn -> Uniq -> Defn
forall a. IntMap a -> Uniq -> a
IM.! Uniq
u