{-# LANGUAGE Strict #-}

-- | Definition of a polymorphic (generic) pass that can work with
-- programs of any rep.
module Futhark.Pass
  ( PassM,
    runPassM,
    Pass (..),
    passLongOption,
    parPass,
    intraproceduralTransformation,
    intraproceduralTransformationWithConsts,
  )
where

import Control.Monad.State.Strict
import Control.Monad.Writer.Strict
import Control.Parallel.Strategies
import Data.Char
import Futhark.IR
import Futhark.MonadFreshNames
import Futhark.Util.Log
import Prelude hiding (log)

-- | The monad in which passes execute.
newtype PassM a = PassM (WriterT Log (State VNameSource) a)
  deriving (forall a b. a -> PassM b -> PassM a
forall a b. (a -> b) -> PassM a -> PassM b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> PassM b -> PassM a
$c<$ :: forall a b. a -> PassM b -> PassM a
fmap :: forall a b. (a -> b) -> PassM a -> PassM b
$cfmap :: forall a b. (a -> b) -> PassM a -> PassM b
Functor, Functor PassM
forall a. a -> PassM a
forall a b. PassM a -> PassM b -> PassM a
forall a b. PassM a -> PassM b -> PassM b
forall a b. PassM (a -> b) -> PassM a -> PassM b
forall a b c. (a -> b -> c) -> PassM a -> PassM b -> PassM c
forall (f :: * -> *).
Functor f
-> (forall a. a -> f a)
-> (forall a b. f (a -> b) -> f a -> f b)
-> (forall a b c. (a -> b -> c) -> f a -> f b -> f c)
-> (forall a b. f a -> f b -> f b)
-> (forall a b. f a -> f b -> f a)
-> Applicative f
<* :: forall a b. PassM a -> PassM b -> PassM a
$c<* :: forall a b. PassM a -> PassM b -> PassM a
*> :: forall a b. PassM a -> PassM b -> PassM b
$c*> :: forall a b. PassM a -> PassM b -> PassM b
liftA2 :: forall a b c. (a -> b -> c) -> PassM a -> PassM b -> PassM c
$cliftA2 :: forall a b c. (a -> b -> c) -> PassM a -> PassM b -> PassM c
<*> :: forall a b. PassM (a -> b) -> PassM a -> PassM b
$c<*> :: forall a b. PassM (a -> b) -> PassM a -> PassM b
pure :: forall a. a -> PassM a
$cpure :: forall a. a -> PassM a
Applicative, Applicative PassM
forall a. a -> PassM a
forall a b. PassM a -> PassM b -> PassM b
forall a b. PassM a -> (a -> PassM b) -> PassM b
forall (m :: * -> *).
Applicative m
-> (forall a b. m a -> (a -> m b) -> m b)
-> (forall a b. m a -> m b -> m b)
-> (forall a. a -> m a)
-> Monad m
return :: forall a. a -> PassM a
$creturn :: forall a. a -> PassM a
>> :: forall a b. PassM a -> PassM b -> PassM b
$c>> :: forall a b. PassM a -> PassM b -> PassM b
>>= :: forall a b. PassM a -> (a -> PassM b) -> PassM b
$c>>= :: forall a b. PassM a -> (a -> PassM b) -> PassM b
Monad)

instance MonadLogger PassM where
  addLog :: Log -> PassM ()
addLog = forall a. WriterT Log (State VNameSource) a -> PassM a
PassM forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall w (m :: * -> *). MonadWriter w m => w -> m ()
tell

instance MonadFreshNames PassM where
  putNameSource :: VNameSource -> PassM ()
putNameSource = forall a. WriterT Log (State VNameSource) a -> PassM a
PassM forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s (m :: * -> *). MonadState s m => s -> m ()
put
  getNameSource :: PassM VNameSource
getNameSource = forall a. WriterT Log (State VNameSource) a -> PassM a
PassM forall s (m :: * -> *). MonadState s m => m s
get

-- | Execute a 'PassM' action, yielding logging information and either
-- an error pretty or a result.
runPassM ::
  MonadFreshNames m =>
  PassM a ->
  m (a, Log)
runPassM :: forall (m :: * -> *) a. MonadFreshNames m => PassM a -> m (a, Log)
runPassM (PassM WriterT Log (State VNameSource) a
m) = forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource forall a b. (a -> b) -> a -> b
$ forall s a. State s a -> s -> (a, s)
runState (forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
runWriterT WriterT Log (State VNameSource) a
m)

-- | A compiler pass transforming a 'Prog' of a given rep to a 'Prog'
-- of another rep.
data Pass fromrep torep = Pass
  { -- | Name of the pass.  Keep this short and simple.  It will
    -- be used to automatically generate a command-line option
    -- name via 'passLongOption'.
    forall fromrep torep. Pass fromrep torep -> String
passName :: String,
    -- | A slightly longer description, which will show up in the
    -- command-line help pretty.
    forall fromrep torep. Pass fromrep torep -> String
passDescription :: String,
    forall fromrep torep.
Pass fromrep torep -> Prog fromrep -> PassM (Prog torep)
passFunction :: Prog fromrep -> PassM (Prog torep)
  }

-- | Take the name of the pass, turn spaces into dashes, and make all
-- characters lowercase.
passLongOption :: Pass fromrep torep -> String
passLongOption :: forall fromrep torep. Pass fromrep torep -> String
passLongOption = forall a b. (a -> b) -> [a] -> [b]
map (Char -> Char
spaceToDash forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> Char
toLower) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall fromrep torep. Pass fromrep torep -> String
passName
  where
    spaceToDash :: Char -> Char
spaceToDash Char
' ' = Char
'-'
    spaceToDash Char
c = Char
c

-- | Apply a 'PassM' operation in parallel to multiple elements,
-- joining together the name sources and logs, and propagating any
-- error properly.
parPass :: (a -> PassM b) -> [a] -> PassM [b]
parPass :: forall a b. (a -> PassM b) -> [a] -> PassM [b]
parPass a -> PassM b
f [a]
as = do
  ([b]
x, Log
log) <- forall (m :: * -> *) a.
MonadFreshNames m =>
(VNameSource -> (a, VNameSource)) -> m a
modifyNameSource forall a b. (a -> b) -> a -> b
$ \VNameSource
src ->
    let ([b]
bs, [Log]
logs, [VNameSource]
srcs) = forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 forall a b. (a -> b) -> a -> b
$ forall b a. Strategy b -> (a -> b) -> [a] -> [b]
parMap forall a. Strategy a
rpar (forall {c}.
MonadFreshNames (StateT c Identity) =>
c -> a -> (b, Log, c)
f' VNameSource
src) [a]
as
     in (([b]
bs, forall a. Monoid a => [a] -> a
mconcat [Log]
logs), forall a. Monoid a => [a] -> a
mconcat [VNameSource]
srcs)

  forall (m :: * -> *). MonadLogger m => Log -> m ()
addLog Log
log
  forall (f :: * -> *) a. Applicative f => a -> f a
pure [b]
x
  where
    f' :: c -> a -> (b, Log, c)
f' c
src a
a =
      let ((b
x', Log
log), c
src') = forall s a. State s a -> s -> (a, s)
runState (forall (m :: * -> *) a. MonadFreshNames m => PassM a -> m (a, Log)
runPassM (a -> PassM b
f a
a)) c
src
       in (b
x', Log
log, c
src')

-- | Apply some operation to the top-level constants. Then applies an
-- operation to all the function definitions, which are also given the
-- transformed constants so they can be brought into scope.
-- The function definition transformations are run in parallel (with
-- 'parPass'), since they cannot affect each other.
intraproceduralTransformationWithConsts ::
  (Stms fromrep -> PassM (Stms torep)) ->
  (Stms torep -> FunDef fromrep -> PassM (FunDef torep)) ->
  Prog fromrep ->
  PassM (Prog torep)
intraproceduralTransformationWithConsts :: forall fromrep torep.
(Stms fromrep -> PassM (Stms torep))
-> (Stms torep -> FunDef fromrep -> PassM (FunDef torep))
-> Prog fromrep
-> PassM (Prog torep)
intraproceduralTransformationWithConsts Stms fromrep -> PassM (Stms torep)
ct Stms torep -> FunDef fromrep -> PassM (FunDef torep)
ft Prog fromrep
prog = do
  Stms torep
consts' <- Stms fromrep -> PassM (Stms torep)
ct (forall rep. Prog rep -> Stms rep
progConsts Prog fromrep
prog)
  [FunDef torep]
funs' <- forall a b. (a -> PassM b) -> [a] -> PassM [b]
parPass (Stms torep -> FunDef fromrep -> PassM (FunDef torep)
ft Stms torep
consts') (forall rep. Prog rep -> [FunDef rep]
progFuns Prog fromrep
prog)
  forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ Prog fromrep
prog {progConsts :: Stms torep
progConsts = Stms torep
consts', progFuns :: [FunDef torep]
progFuns = [FunDef torep]
funs'}

-- | Like 'intraproceduralTransformationWithConsts', but do not change
-- the top-level constants, and simply pass along their 'Scope'.
intraproceduralTransformation ::
  (Scope rep -> Stms rep -> PassM (Stms rep)) ->
  Prog rep ->
  PassM (Prog rep)
intraproceduralTransformation :: forall rep.
(Scope rep -> Stms rep -> PassM (Stms rep))
-> Prog rep -> PassM (Prog rep)
intraproceduralTransformation Scope rep -> Stms rep -> PassM (Stms rep)
f =
  forall fromrep torep.
(Stms fromrep -> PassM (Stms torep))
-> (Stms torep -> FunDef fromrep -> PassM (FunDef torep))
-> Prog fromrep
-> PassM (Prog torep)
intraproceduralTransformationWithConsts (Scope rep -> Stms rep -> PassM (Stms rep)
f forall a. Monoid a => a
mempty) forall {a}. Scoped rep a => a -> FunDef rep -> PassM (FunDef rep)
f'
  where
    f' :: a -> FunDef rep -> PassM (FunDef rep)
f' a
consts FunDef rep
fd = do
      Stms rep
stms <-
        Scope rep -> Stms rep -> PassM (Stms rep)
f
          (forall rep a. Scoped rep a => a -> Scope rep
scopeOf a
consts forall a. Semigroup a => a -> a -> a
<> forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
scopeOfFParams (forall rep. FunDef rep -> [FParam rep]
funDefParams FunDef rep
fd))
          (forall rep. Body rep -> Stms rep
bodyStms forall a b. (a -> b) -> a -> b
$ forall rep. FunDef rep -> Body rep
funDefBody FunDef rep
fd)
      forall (f :: * -> *) a. Applicative f => a -> f a
pure FunDef rep
fd {funDefBody :: Body rep
funDefBody = (forall rep. FunDef rep -> Body rep
funDefBody FunDef rep
fd) {bodyStms :: Stms rep
bodyStms = Stms rep
stms}}