    {-# 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 -> b) -> PassM a -> PassM b)
    -> (forall a b. a -> PassM b -> PassM a) -> Functor PassM
    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
    $cfmap :: forall a b. (a -> b) -> PassM a -> PassM b
    fmap :: forall a b. (a -> b) -> PassM a -> PassM b
    $c<$ :: forall a b. a -> PassM b -> PassM a
    <$ :: forall a b. a -> PassM b -> PassM a
    Functor, Functor PassM
    Functor PassM =>
    (forall a. a -> PassM a)
    -> (forall a b. PassM (a -> b) -> PassM a -> PassM b)
    -> (forall a b c. (a -> b -> c) -> PassM a -> PassM b -> PassM c)
    -> (forall a b. PassM a -> PassM b -> PassM b)
    -> (forall a b. PassM a -> PassM b -> PassM a)
    -> Applicative 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
    $cpure :: forall a. a -> PassM a
    pure :: forall a. a -> PassM a
    $c<*> :: forall a b. PassM (a -> b) -> PassM a -> PassM b
    <*> :: forall a b. PassM (a -> b) -> PassM a -> PassM b
    $cliftA2 :: forall a b c. (a -> b -> c) -> PassM a -> PassM b -> PassM c
    liftA2 :: forall a b c. (a -> b -> c) -> PassM a -> PassM b -> PassM c
    $c*> :: forall a b. PassM a -> PassM b -> PassM b
    *> :: forall a b. PassM a -> PassM b -> PassM b
    $c<* :: forall a b. PassM a -> PassM b -> PassM a
    <* :: forall a b. PassM a -> PassM b -> PassM a
    Applicative, Applicative PassM
    Applicative PassM =>
    (forall a b. PassM a -> (a -> PassM b) -> PassM b)
    -> (forall a b. PassM a -> PassM b -> PassM b)
    -> (forall a. a -> PassM a)
    -> Monad 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
    $c>>= :: forall a b. PassM a -> (a -> PassM b) -> PassM b
    >>= :: forall a b. PassM a -> (a -> PassM b) -> PassM b
    $c>> :: forall a b. PassM a -> PassM b -> PassM b
    >> :: forall a b. PassM a -> PassM b -> PassM b
    $creturn :: forall a. a -> PassM a
    return :: forall a. a -> PassM a
    Monad)

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

    instance MonadFreshNames PassM where
      putNameSource :: VNameSource -> PassM ()
    putNameSource = WriterT Log (State VNameSource) () -> PassM ()
    forall a. WriterT Log (State VNameSource) a -> PassM a
    PassM (WriterT Log (State VNameSource) () -> PassM ())
    -> (VNameSource -> WriterT Log (State VNameSource) ())
    -> VNameSource
    -> PassM ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . VNameSource -> WriterT Log (State VNameSource) ()
    forall s (m :: * -> *). MonadState s m => s -> m ()
    put
      getNameSource :: PassM VNameSource
    getNameSource = WriterT Log (State VNameSource) VNameSource -> PassM VNameSource
    forall a. WriterT Log (State VNameSource) a -> PassM a
    PassM WriterT Log (State VNameSource) VNameSource
    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) = (VNameSource -> ((a, Log), VNameSource)) -> m (a, Log)
    forall (m :: * -> *) a.
    MonadFreshNames m =>
    (VNameSource -> (a, VNameSource)) -> m a
    modifyNameSource ((VNameSource -> ((a, Log), VNameSource)) -> m (a, Log))
    -> (VNameSource -> ((a, Log), VNameSource)) -> m (a, Log)
    forall a b. (a -> b) -> a -> b
    $ State VNameSource (a, Log)
    -> VNameSource -> ((a, Log), VNameSource)
    forall s a. State s a -> s -> (a, s)
    runState (WriterT Log (State VNameSource) a -> State VNameSource (a, Log)
    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 option.
        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 = (Char -> Char) -> String -> String
    forall a b. (a -> b) -> [a] -> [b]
    map (Char -> Char
    spaceToDash (Char -> Char) -> (Char -> Char) -> Char -> Char
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Char -> Char
    toLower) (String -> String)
    -> (Pass fromrep torep -> String) -> Pass fromrep torep -> String
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Pass fromrep torep -> String
    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) <- (VNameSource -> (([b], Log), VNameSource)) -> PassM ([b], Log)
    forall (m :: * -> *) a.
    MonadFreshNames m =>
    (VNameSource -> (a, VNameSource)) -> m a
    modifyNameSource ((VNameSource -> (([b], Log), VNameSource)) -> PassM ([b], Log))
    -> (VNameSource -> (([b], Log), VNameSource)) -> PassM ([b], Log)
    forall a b. (a -> b) -> a -> b
    $ \VNameSource
    src ->
        let ([b]
    bs, [Log]
    logs, [VNameSource]
    srcs) = [(b, Log, VNameSource)] -> ([b], [Log], [VNameSource])
    forall a b c. [(a, b, c)] -> ([a], [b], [c])
    unzip3 ([(b, Log, VNameSource)] -> ([b], [Log], [VNameSource]))
    -> [(b, Log, VNameSource)] -> ([b], [Log], [VNameSource])
    forall a b. (a -> b) -> a -> b
    $ Strategy (b, Log, VNameSource)
    -> (a -> (b, Log, VNameSource)) -> [a] -> [(b, Log, VNameSource)]
    forall b a. Strategy b -> (a -> b) -> [a] -> [b]
    parMap Strategy (b, Log, VNameSource)
    forall a. Strategy a
    rpar (VNameSource -> a -> (b, Log, VNameSource)
    forall {c}.
    MonadFreshNames (StateT c Identity) =>
    c -> a -> (b, Log, c)
    f' VNameSource
    src) [a]
    as
         in (([b]
    bs, [Log] -> Log
    forall a. Monoid a => [a] -> a
    mconcat [Log]
    logs), [VNameSource] -> VNameSource
    forall a. Monoid a => [a] -> a
    mconcat [VNameSource]
    srcs)

      Log -> PassM ()
    forall (m :: * -> *). MonadLogger m => Log -> m ()
    addLog Log
    log
      [b] -> PassM [b]
    forall a. a -> PassM a
    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') = State c (b, Log) -> c -> ((b, Log), c)
    forall s a. State s a -> s -> (a, s)
    runState (PassM b -> State c (b, Log)
    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 (Prog fromrep -> Stms fromrep
    forall rep. Prog rep -> Stms rep
    progConsts Prog fromrep
    prog)
      [FunDef torep]
    funs' <- (FunDef fromrep -> PassM (FunDef torep))
    -> [FunDef fromrep] -> PassM [FunDef torep]
    forall a b. (a -> PassM b) -> [a] -> PassM [b]
    parPass (Stms torep -> FunDef fromrep -> PassM (FunDef torep)
    ft Stms torep
    consts') (Prog fromrep -> [FunDef fromrep]
    forall rep. Prog rep -> [FunDef rep]
    progFuns Prog fromrep
    prog)
      Prog torep -> PassM (Prog torep)
    forall a. a -> PassM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Prog torep -> PassM (Prog torep))
    -> Prog torep -> PassM (Prog torep)
    forall a b. (a -> b) -> a -> b
    $ Prog fromrep
    prog {progConsts = consts', progFuns = 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 =
      (Stms rep -> PassM (Stms rep))
    -> (Stms rep -> FunDef rep -> PassM (FunDef rep))
    -> Prog rep
    -> PassM (Prog rep)
    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 Scope rep
    forall a. Monoid a => a
    mempty) Stms rep -> FunDef rep -> PassM (FunDef rep)
    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
              (a -> Scope rep
    forall rep a. Scoped rep a => a -> Scope rep
    scopeOf a
    consts Scope rep -> Scope rep -> Scope rep
    forall a. Semigroup a => a -> a -> a
    <> [Param (FParamInfo rep)] -> Scope rep
    forall rep dec. (FParamInfo rep ~ dec) => [Param dec] -> Scope rep
    scopeOfFParams (FunDef rep -> [Param (FParamInfo rep)]
    forall rep. FunDef rep -> [FParam rep]
    funDefParams FunDef rep
    fd))
              (Body rep -> Stms rep
    forall rep. Body rep -> Stms rep
    bodyStms (Body rep -> Stms rep) -> Body rep -> Stms rep
    forall a b. (a -> b) -> a -> b
    $ FunDef rep -> Body rep
    forall rep. FunDef rep -> Body rep
    funDefBody FunDef rep
    fd)
          FunDef rep -> PassM (FunDef rep)
    forall a. a -> PassM a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure FunDef rep
    fd {funDefBody = (funDefBody fd) {bodyStms = stms}}