    {-# LANGUAGE UndecidableInstances #-}

    -- | This module implements common-subexpression elimination.  This
    -- module does not actually remove the duplicate, but only replaces
    -- one with a diference to the other.  E.g:
    --
    -- @
    --   let a = x + y
    --   let b = x + y
    -- @
    --
    -- becomes:
    --
    -- @
    --   let a = x + y
    --   let b = a
    -- @
    --
    -- After which copy propagation in the simplifier will actually remove
    -- the definition of @b@.
    --
    -- Our CSE is still rather stupid.  No normalisation is performed, so
    -- the expressions @x+y@ and @y+x@ will be considered distinct.
    -- Furthermore, no expression with its own binding will be considered
    -- equal to any other, since the variable names will be distinct.
    -- This affects SOACs in particular.
    module Futhark.Optimise.CSE
      ( performCSE,
        performCSEOnFunDef,
        performCSEOnStms,
        CSEInOp,
      )
    where

    import Control.Monad.Reader
    import Data.Map.Strict qualified as M
    import Futhark.Analysis.Alias
    import Futhark.IR
    import Futhark.IR.Aliases
      ( Aliases,
        consumedInStms,
        mkStmsAliases,
        removeFunDefAliases,
        removeProgAliases,
        removeStmAliases,
      )
    import Futhark.IR.GPU qualified as GPU
    import Futhark.IR.MC qualified as MC
    import Futhark.IR.Mem qualified as Memory
    import Futhark.IR.Prop.Aliases
    import Futhark.IR.SOACS.SOAC qualified as SOAC
    import Futhark.Pass
    import Futhark.Transform.Substitute

    -- | Perform CSE on every function in a program.
    --
    -- If the boolean argument is false, the pass will not perform CSE on
    -- expressions producing arrays. This should be disabled when the rep has
    -- memory information, since at that point arrays have identity beyond their
    -- value.
    performCSE ::
      (AliasableRep rep, CSEInOp (Op (Aliases rep))) =>
      Bool ->
      Pass rep rep
    performCSE :: forall rep.
    (AliasableRep rep, CSEInOp (Op (Aliases rep))) =>
    Bool -> Pass rep rep
    performCSE Bool
    cse_arrays =
      String -> String -> (Prog rep -> PassM (Prog rep)) -> Pass rep rep
    forall fromrep torep.
    String
    -> String
    -> (Prog fromrep -> PassM (Prog torep))
    -> Pass fromrep torep
    Pass String
    "CSE" String
    "Combine common subexpressions." ((Prog rep -> PassM (Prog rep)) -> Pass rep rep)
    -> (Prog rep -> PassM (Prog rep)) -> Pass rep rep
    forall a b. (a -> b) -> a -> b
    $ \Prog rep
    prog ->
        (Prog (Aliases rep) -> Prog rep)
    -> PassM (Prog (Aliases rep)) -> PassM (Prog rep)
    forall a b. (a -> b) -> PassM a -> PassM b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap Prog (Aliases rep) -> Prog rep
    forall rep. RephraseOp (OpC rep) => Prog (Aliases rep) -> Prog rep
    removeProgAliases
          (PassM (Prog (Aliases rep)) -> PassM (Prog rep))
    -> (Prog rep -> PassM (Prog (Aliases rep)))
    -> Prog rep
    -> PassM (Prog rep)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Stms (Aliases rep) -> PassM (Stms (Aliases rep)))
    -> (Stms (Aliases rep)
        -> FunDef (Aliases rep) -> PassM (FunDef (Aliases rep)))
    -> Prog (Aliases rep)
    -> PassM (Prog (Aliases rep))
    forall fromrep torep.
    (Stms fromrep -> PassM (Stms torep))
    -> (Stms torep -> FunDef fromrep -> PassM (FunDef torep))
    -> Prog fromrep
    -> PassM (Prog torep)
    intraproceduralTransformationWithConsts
            (Names -> Stms (Aliases rep) -> PassM (Stms (Aliases rep))
    forall {f :: * -> *} {rep}.
    (Applicative f, CSEInOp (OpC rep rep), Aliased rep) =>
    Names -> Seq (Stm rep) -> f (Seq (Stm rep))
    onConsts ([FunDef rep] -> Names
    forall a. FreeIn a => a -> Names
    freeIn (Prog rep -> [FunDef rep]
    forall rep. Prog rep -> [FunDef rep]
    progFuns Prog rep
    prog)))
            Stms (Aliases rep)
    -> FunDef (Aliases rep) -> PassM (FunDef (Aliases rep))
    forall {f :: * -> *} {rep} {p}.
    (Applicative f, Aliased rep, CSEInOp (OpC rep rep)) =>
    p -> FunDef rep -> f (FunDef rep)
    onFun
          (Prog (Aliases rep) -> PassM (Prog (Aliases rep)))
    -> (Prog rep -> Prog (Aliases rep))
    -> Prog rep
    -> PassM (Prog (Aliases rep))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Prog rep -> Prog (Aliases rep)
    forall rep. AliasableRep rep => Prog rep -> Prog (Aliases rep)
    aliasAnalysis
          (Prog rep -> PassM (Prog rep)) -> Prog rep -> PassM (Prog rep)
    forall a b. (a -> b) -> a -> b
    $ Prog rep
    prog
      where
        onConsts :: Names -> Seq (Stm rep) -> f (Seq (Stm rep))
    onConsts Names
    free_in_funs Seq (Stm rep)
    stms = do
          let free_list :: [VName]
    free_list = Names -> [VName]
    namesToList Names
    free_in_funs
              ([Names]
    res_als, Names
    stms_cons) = Seq (Stm rep) -> Result -> ([Names], Names)
    forall rep. Aliased rep => Stms rep -> Result -> ([Names], Names)
    mkStmsAliases Seq (Stm rep)
    stms (Result -> ([Names], Names)) -> Result -> ([Names], Names)
    forall a b. (a -> b) -> a -> b
    $ [VName] -> Result
    varsRes [VName]
    free_list
          Seq (Stm rep) -> f (Seq (Stm rep))
    forall a. a -> f a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Seq (Stm rep) -> f (Seq (Stm rep)))
    -> ((Seq (Stm rep), ()) -> Seq (Stm rep))
    -> (Seq (Stm rep), ())
    -> f (Seq (Stm rep))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Seq (Stm rep), ()) -> Seq (Stm rep)
    forall a b. (a, b) -> a
    fst ((Seq (Stm rep), ()) -> f (Seq (Stm rep)))
    -> (Seq (Stm rep), ()) -> f (Seq (Stm rep))
    forall a b. (a -> b) -> a -> b
    $
            Reader (CSEState rep) (Seq (Stm rep), ())
    -> CSEState rep -> (Seq (Stm rep), ())
    forall r a. Reader r a -> r -> a
    runReader
              ( Names
    -> [Stm rep]
    -> CSEM rep ()
    -> Reader (CSEState rep) (Seq (Stm rep), ())
    forall rep a.
    (Aliased rep, CSEInOp (Op rep)) =>
    Names -> [Stm rep] -> CSEM rep a -> CSEM rep (Stms rep, a)
    cseInStms
                  ([Names] -> Names
    forall a. Monoid a => [a] -> a
    mconcat [Names]
    res_als Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> Names
    stms_cons)
                  (Seq (Stm rep) -> [Stm rep]
    forall rep. Stms rep -> [Stm rep]
    stmsToList Seq (Stm rep)
    stms)
                  (() -> CSEM rep ()
    forall a. a -> ReaderT (CSEState rep) Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ())
              )
              (Bool -> CSEState rep
    forall rep. Bool -> CSEState rep
    newCSEState Bool
    cse_arrays)
        onFun :: p -> FunDef rep -> f (FunDef rep)
    onFun p
    _ = FunDef rep -> f (FunDef rep)
    forall a. a -> f a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (FunDef rep -> f (FunDef rep))
    -> (FunDef rep -> FunDef rep) -> FunDef rep -> f (FunDef rep)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Bool -> FunDef rep -> FunDef rep
    forall rep.
    (Aliased rep, CSEInOp (Op rep)) =>
    Bool -> FunDef rep -> FunDef rep
    cseInFunDef Bool
    cse_arrays

    -- | Perform CSE on a single function.
    --
    -- If the boolean argument is false, the pass will not perform CSE on
    -- expressions producing arrays. This should be disabled when the rep has
    -- memory information, since at that point arrays have identity beyond their
    -- value.
    performCSEOnFunDef ::
      (AliasableRep rep, CSEInOp (Op (Aliases rep))) =>
      Bool ->
      FunDef rep ->
      FunDef rep
    performCSEOnFunDef :: forall rep.
    (AliasableRep rep, CSEInOp (Op (Aliases rep))) =>
    Bool -> FunDef rep -> FunDef rep
    performCSEOnFunDef Bool
    cse_arrays =
      FunDef (Aliases rep) -> FunDef rep
    forall rep.
    RephraseOp (OpC rep) =>
    FunDef (Aliases rep) -> FunDef rep
    removeFunDefAliases (FunDef (Aliases rep) -> FunDef rep)
    -> (FunDef rep -> FunDef (Aliases rep)) -> FunDef rep -> FunDef rep
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Bool -> FunDef (Aliases rep) -> FunDef (Aliases rep)
    forall rep.
    (Aliased rep, CSEInOp (Op rep)) =>
    Bool -> FunDef rep -> FunDef rep
    cseInFunDef Bool
    cse_arrays (FunDef (Aliases rep) -> FunDef (Aliases rep))
    -> (FunDef rep -> FunDef (Aliases rep))
    -> FunDef rep
    -> FunDef (Aliases rep)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . FunDef rep -> FunDef (Aliases rep)
    forall rep. AliasableRep rep => FunDef rep -> FunDef (Aliases rep)
    analyseFun

    -- | Perform CSE on some statements.
    performCSEOnStms ::
      (AliasableRep rep, CSEInOp (Op (Aliases rep))) =>
      Stms rep ->
      Stms rep
    performCSEOnStms :: forall rep.
    (AliasableRep rep, CSEInOp (Op (Aliases rep))) =>
    Stms rep -> Stms rep
    performCSEOnStms =
      (Stm (Aliases rep) -> Stm rep)
    -> Seq (Stm (Aliases rep)) -> Seq (Stm rep)
    forall a b. (a -> b) -> Seq a -> Seq b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap Stm (Aliases rep) -> Stm rep
    forall rep. RephraseOp (OpC rep) => Stm (Aliases rep) -> Stm rep
    removeStmAliases (Seq (Stm (Aliases rep)) -> Seq (Stm rep))
    -> (Seq (Stm rep) -> Seq (Stm (Aliases rep)))
    -> Seq (Stm rep)
    -> Seq (Stm rep)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Seq (Stm (Aliases rep)) -> Seq (Stm (Aliases rep))
    forall {rep}.
    (Aliased rep, CSEInOp (OpC rep rep)) =>
    Stms rep -> Stms rep
    f (Seq (Stm (Aliases rep)) -> Seq (Stm (Aliases rep)))
    -> (Seq (Stm rep) -> Seq (Stm (Aliases rep)))
    -> Seq (Stm rep)
    -> Seq (Stm (Aliases rep))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Seq (Stm (Aliases rep)), AliasesAndConsumed)
    -> Seq (Stm (Aliases rep))
    forall a b. (a, b) -> a
    fst ((Seq (Stm (Aliases rep)), AliasesAndConsumed)
     -> Seq (Stm (Aliases rep)))
    -> (Seq (Stm rep) -> (Seq (Stm (Aliases rep)), AliasesAndConsumed))
    -> Seq (Stm rep)
    -> Seq (Stm (Aliases rep))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . AliasTable
    -> Seq (Stm rep) -> (Seq (Stm (Aliases rep)), AliasesAndConsumed)
    forall rep.
    AliasableRep rep =>
    AliasTable -> Stms rep -> (Stms (Aliases rep), AliasesAndConsumed)
    analyseStms AliasTable
    forall a. Monoid a => a
    mempty
      where
        f :: Stms rep -> Stms rep
    f Stms rep
    stms =
          (Stms rep, ()) -> Stms rep
    forall a b. (a, b) -> a
    fst ((Stms rep, ()) -> Stms rep) -> (Stms rep, ()) -> Stms rep
    forall a b. (a -> b) -> a -> b
    $
            Reader (CSEState rep) (Stms rep, ())
    -> CSEState rep -> (Stms rep, ())
    forall r a. Reader r a -> r -> a
    runReader
              (Names
    -> [Stm rep] -> CSEM rep () -> Reader (CSEState rep) (Stms rep, ())
    forall rep a.
    (Aliased rep, CSEInOp (Op rep)) =>
    Names -> [Stm rep] -> CSEM rep a -> CSEM rep (Stms rep, a)
    cseInStms (Stms rep -> Names
    forall rep. Aliased rep => Stms rep -> Names
    consumedInStms Stms rep
    stms) (Stms rep -> [Stm rep]
    forall rep. Stms rep -> [Stm rep]
    stmsToList Stms rep
    stms) (() -> CSEM rep ()
    forall a. a -> ReaderT (CSEState rep) Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ()))
              -- It is never safe to CSE arrays in stms in isolation,
              -- because we might introduce additional aliasing.
              (Bool -> CSEState rep
    forall rep. Bool -> CSEState rep
    newCSEState Bool
    False)

    cseInFunDef ::
      (Aliased rep, CSEInOp (Op rep)) =>
      Bool ->
      FunDef rep ->
      FunDef rep
    cseInFunDef :: forall rep.
    (Aliased rep, CSEInOp (Op rep)) =>
    Bool -> FunDef rep -> FunDef rep
    cseInFunDef Bool
    cse_arrays FunDef rep
    fundec =
      FunDef rep
    fundec
        { funDefBody =
            runReader (cseInBody ds $ funDefBody fundec) $ newCSEState cse_arrays
        }
      where
        -- XXX: we treat every array result as a consumption here, because
        -- it is otherwise complicated to ensure we do not introduce more
        -- aliasing than specified by the return type. This is not a
        -- practical problem while we still perform such aggressive
        -- inlining.
        ds :: [Diet]
    ds = ((RetType rep, RetAls) -> Diet)
    -> [(RetType rep, RetAls)] -> [Diet]
    forall a b. (a -> b) -> [a] -> [b]
    map (RetType rep -> Diet
    forall {t}. DeclExtTyped t => t -> Diet
    retDiet (RetType rep -> Diet)
    -> ((RetType rep, RetAls) -> RetType rep)
    -> (RetType rep, RetAls)
    -> Diet
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (RetType rep, RetAls) -> RetType rep
    forall a b. (a, b) -> a
    fst) ([(RetType rep, RetAls)] -> [Diet])
    -> [(RetType rep, RetAls)] -> [Diet]
    forall a b. (a -> b) -> a -> b
    $ FunDef rep -> [(RetType rep, RetAls)]
    forall rep. FunDef rep -> [(RetType rep, RetAls)]
    funDefRetType FunDef rep
    fundec
        retDiet :: t -> Diet
    retDiet t
    t
          | TypeBase ExtShape Uniqueness -> Bool
    forall shape u. TypeBase shape u -> Bool
    primType (TypeBase ExtShape Uniqueness -> Bool)
    -> TypeBase ExtShape Uniqueness -> Bool
    forall a b. (a -> b) -> a -> b
    $ t -> TypeBase ExtShape Uniqueness
    forall t. DeclExtTyped t => t -> TypeBase ExtShape Uniqueness
    declExtTypeOf t
    t = Diet
    Observe
          | Bool
    otherwise = Diet
    Consume

    type CSEM rep = Reader (CSEState rep)

    cseInBody ::
      (Aliased rep, CSEInOp (Op rep)) =>
      [Diet] ->
      Body rep ->
      CSEM rep (Body rep)
    cseInBody :: forall rep.
    (Aliased rep, CSEInOp (Op rep)) =>
    [Diet] -> Body rep -> CSEM rep (Body rep)
    cseInBody [Diet]
    ds (Body BodyDec rep
    bodydec Stms rep
    stms Result
    res) = do
      (Stms rep
    stms', Result
    res') <-
        Names
    -> [Stm rep] -> CSEM rep Result -> CSEM rep (Stms rep, Result)
    forall rep a.
    (Aliased rep, CSEInOp (Op rep)) =>
    Names -> [Stm rep] -> CSEM rep a -> CSEM rep (Stms rep, a)
    cseInStms (Names
    res_cons Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> Names
    stms_cons) (Stms rep -> [Stm rep]
    forall rep. Stms rep -> [Stm rep]
    stmsToList Stms rep
    stms) (CSEM rep Result -> CSEM rep (Stms rep, Result))
    -> CSEM rep Result -> CSEM rep (Stms rep, Result)
    forall a b. (a -> b) -> a -> b
    $ do
          CSEState (ExpressionSubstitutions rep
    _, NameSubstitutions
    nsubsts) Bool
    _ <- ReaderT (CSEState rep) Identity (CSEState rep)
    forall r (m :: * -> *). MonadReader r m => m r
    ask
          Result -> CSEM rep Result
    forall a. a -> ReaderT (CSEState rep) Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Result -> CSEM rep Result) -> Result -> CSEM rep Result
    forall a b. (a -> b) -> a -> b
    $ NameSubstitutions -> Result -> Result
    forall a. Substitute a => NameSubstitutions -> a -> a
    substituteNames NameSubstitutions
    nsubsts Result
    res
      Body rep -> CSEM rep (Body rep)
    forall a. a -> ReaderT (CSEState rep) Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Body rep -> CSEM rep (Body rep))
    -> Body rep -> CSEM rep (Body rep)
    forall a b. (a -> b) -> a -> b
    $ BodyDec rep -> Stms rep -> Result -> Body rep
    forall rep. BodyDec rep -> Stms rep -> Result -> Body rep
    Body BodyDec rep
    bodydec Stms rep
    stms' Result
    res'
      where
        ([Names]
    res_als, Names
    stms_cons) = Stms rep -> Result -> ([Names], Names)
    forall rep. Aliased rep => Stms rep -> Result -> ([Names], Names)
    mkStmsAliases Stms rep
    stms Result
    res
        res_cons :: Names
    res_cons = [Names] -> Names
    forall a. Monoid a => [a] -> a
    mconcat ([Names] -> Names) -> [Names] -> Names
    forall a b. (a -> b) -> a -> b
    $ (Diet -> Names -> Names) -> [Diet] -> [Names] -> [Names]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith Diet -> Names -> Names
    forall {p}. Monoid p => Diet -> p -> p
    consumeResult [Diet]
    ds [Names]
    res_als
        consumeResult :: Diet -> p -> p
    consumeResult Diet
    Consume p
    als = p
    als
        consumeResult Diet
    _ p
    _ = p
    forall a. Monoid a => a
    mempty

    cseInLambda ::
      (Aliased rep, CSEInOp (Op rep)) =>
      Lambda rep ->
      CSEM rep (Lambda rep)
    cseInLambda :: forall rep.
    (Aliased rep, CSEInOp (Op rep)) =>
    Lambda rep -> CSEM rep (Lambda rep)
    cseInLambda Lambda rep
    lam = do
      Body rep
    body' <- [Diet] -> Body rep -> CSEM rep (Body rep)
    forall rep.
    (Aliased rep, CSEInOp (Op rep)) =>
    [Diet] -> Body rep -> CSEM rep (Body rep)
    cseInBody ((Type -> Diet) -> [Type] -> [Diet]
    forall a b. (a -> b) -> [a] -> [b]
    map (Diet -> Type -> Diet
    forall a b. a -> b -> a
    const Diet
    Observe) ([Type] -> [Diet]) -> [Type] -> [Diet]
    forall a b. (a -> b) -> a -> b
    $ Lambda rep -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda rep
    lam) (Body rep -> CSEM rep (Body rep))
    -> Body rep -> CSEM rep (Body rep)
    forall a b. (a -> b) -> a -> b
    $ Lambda rep -> Body rep
    forall rep. Lambda rep -> Body rep
    lambdaBody Lambda rep
    lam
      Lambda rep -> CSEM rep (Lambda rep)
    forall a. a -> ReaderT (CSEState rep) Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Lambda rep
    lam {lambdaBody = body'}

    cseInStms ::
      forall rep a.
      (Aliased rep, CSEInOp (Op rep)) =>
      Names ->
      [Stm rep] ->
      CSEM rep a ->
      CSEM rep (Stms rep, a)
    cseInStms :: forall rep a.
    (Aliased rep, CSEInOp (Op rep)) =>
    Names -> [Stm rep] -> CSEM rep a -> CSEM rep (Stms rep, a)
    cseInStms Names
    _ [] CSEM rep a
    m = do
      a
    a <- CSEM rep a
    m
      (Stms rep, a) -> CSEM rep (Stms rep, a)
    forall a. a -> ReaderT (CSEState rep) Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Stms rep
    forall a. Monoid a => a
    mempty, a
    a)
    cseInStms Names
    consumed (Stm rep
    stm : [Stm rep]
    stms) CSEM rep a
    m =
      Names
    -> Stm rep
    -> ([Stm rep] -> CSEM rep (Stms rep, a))
    -> CSEM rep (Stms rep, a)
    forall rep a.
    ASTRep rep =>
    Names -> Stm rep -> ([Stm rep] -> CSEM rep a) -> CSEM rep a
    cseInStm Names
    consumed Stm rep
    stm (([Stm rep] -> CSEM rep (Stms rep, a)) -> CSEM rep (Stms rep, a))
    -> ([Stm rep] -> CSEM rep (Stms rep, a)) -> CSEM rep (Stms rep, a)
    forall a b. (a -> b) -> a -> b
    $ \[Stm rep]
    stm' -> do
        (Stms rep
    stms', a
    a) <- Names -> [Stm rep] -> CSEM rep a -> CSEM rep (Stms rep, a)
    forall rep a.
    (Aliased rep, CSEInOp (Op rep)) =>
    Names -> [Stm rep] -> CSEM rep a -> CSEM rep (Stms rep, a)
    cseInStms Names
    consumed [Stm rep]
    stms CSEM rep a
    m
        [Stm rep]
    stm'' <- (Stm rep -> ReaderT (CSEState rep) Identity (Stm rep))
    -> [Stm rep] -> ReaderT (CSEState rep) Identity [Stm rep]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM Stm rep -> ReaderT (CSEState rep) Identity (Stm rep)
    forall {rep}.
    (FParamInfo rep ~ FParamInfo rep, RetType rep ~ RetType rep,
     BranchType rep ~ BranchType rep, LParamInfo rep ~ LParamInfo rep,
     Aliased rep, CSEInOp (OpC rep rep)) =>
    Stm rep -> ReaderT (CSEState rep) Identity (Stm rep)
    nestedCSE [Stm rep]
    stm'
        (Stms rep, a) -> CSEM rep (Stms rep, a)
    forall a. a -> ReaderT (CSEState rep) Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ([Stm rep] -> Stms rep
    forall rep. [Stm rep] -> Stms rep
    stmsFromList [Stm rep]
    stm'' Stms rep -> Stms rep -> Stms rep
    forall a. Semigroup a => a -> a -> a
    <> Stms rep
    stms', a
    a)
      where
        nestedCSE :: Stm rep -> ReaderT (CSEState rep) Identity (Stm rep)
    nestedCSE Stm rep
    stm' = do
          let ds :: [Diet]
    ds =
                case Stm rep -> Exp rep
    forall rep. Stm rep -> Exp rep
    stmExp Stm rep
    stm' of
                  Loop [(Param (FParamInfo rep), SubExp)]
    merge LoopForm
    _ Body rep
    _ -> ((Param (FParamInfo rep), SubExp) -> Diet)
    -> [(Param (FParamInfo rep), SubExp)] -> [Diet]
    forall a b. (a -> b) -> [a] -> [b]
    map (TypeBase Shape Uniqueness -> Diet
    forall shape. TypeBase shape Uniqueness -> Diet
    diet (TypeBase Shape Uniqueness -> Diet)
    -> ((Param (FParamInfo rep), SubExp) -> TypeBase Shape Uniqueness)
    -> (Param (FParamInfo rep), SubExp)
    -> Diet
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . Param (FParamInfo rep) -> TypeBase Shape Uniqueness
    forall t. DeclTyped t => t -> TypeBase Shape Uniqueness
    declTypeOf (Param (FParamInfo rep) -> TypeBase Shape Uniqueness)
    -> ((Param (FParamInfo rep), SubExp) -> Param (FParamInfo rep))
    -> (Param (FParamInfo rep), SubExp)
    -> TypeBase Shape Uniqueness
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Param (FParamInfo rep), SubExp) -> Param (FParamInfo rep)
    forall a b. (a, b) -> a
    fst) [(Param (FParamInfo rep), SubExp)]
    merge
                  Exp rep
    _ -> (PatElem (LetDec rep) -> Diet) -> [PatElem (LetDec rep)] -> [Diet]
    forall a b. (a -> b) -> [a] -> [b]
    map PatElem (LetDec rep) -> Diet
    forall {dec}. PatElem dec -> Diet
    patElemDiet ([PatElem (LetDec rep)] -> [Diet])
    -> [PatElem (LetDec rep)] -> [Diet]
    forall a b. (a -> b) -> a -> b
    $ Pat (LetDec rep) -> [PatElem (LetDec rep)]
    forall dec. Pat dec -> [PatElem dec]
    patElems (Pat (LetDec rep) -> [PatElem (LetDec rep)])
    -> Pat (LetDec rep) -> [PatElem (LetDec rep)]
    forall a b. (a -> b) -> a -> b
    $ Stm rep -> Pat (LetDec rep)
    forall rep. Stm rep -> Pat (LetDec rep)
    stmPat Stm rep
    stm'
          Exp rep
    e <- Mapper rep rep (ReaderT (CSEState rep) Identity)
    -> Exp rep -> ReaderT (CSEState rep) Identity (Exp rep)
    forall (m :: * -> *) frep trep.
    Monad m =>
    Mapper frep trep m -> Exp frep -> m (Exp trep)
    mapExpM ([Diet] -> Mapper rep rep (ReaderT (CSEState rep) Identity)
    forall {frep}.
    (RetType frep ~ RetType rep, BranchType frep ~ BranchType rep,
     FParamInfo frep ~ FParamInfo rep, LParamInfo frep ~ LParamInfo rep,
     Aliased frep, CSEInOp (OpC frep frep)) =>
    [Diet] -> Mapper frep frep (ReaderT (CSEState frep) Identity)
    cse [Diet]
    ds) (Exp rep -> ReaderT (CSEState rep) Identity (Exp rep))
    -> Exp rep -> ReaderT (CSEState rep) Identity (Exp rep)
    forall a b. (a -> b) -> a -> b
    $ Stm rep -> Exp rep
    forall rep. Stm rep -> Exp rep
    stmExp Stm rep
    stm'
          Stm rep -> ReaderT (CSEState rep) Identity (Stm rep)
    forall a. a -> ReaderT (CSEState rep) Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Stm rep
    stm' {stmExp = e}

        cse :: [Diet] -> Mapper frep frep (ReaderT (CSEState frep) Identity)
    cse [Diet]
    ds =
          (forall rep (m :: * -> *). Monad m => Mapper rep rep m
    identityMapper @rep)
            { mapOnBody = const $ cseInBody ds,
              mapOnOp = cseInOp
            }

        patElemDiet :: PatElem dec -> Diet
    patElemDiet PatElem dec
    pe
          | PatElem dec -> VName
    forall dec. PatElem dec -> VName
    patElemName PatElem dec
    pe VName -> Names -> Bool
    `nameIn` Names
    consumed = Diet
    Consume
          | Bool
    otherwise = Diet
    Observe

    cseInStm ::
      (ASTRep rep) =>
      Names ->
      Stm rep ->
      ([Stm rep] -> CSEM rep a) ->
      CSEM rep a
    cseInStm :: forall rep a.
    ASTRep rep =>
    Names -> Stm rep -> ([Stm rep] -> CSEM rep a) -> CSEM rep a
    cseInStm Names
    consumed (Let Pat (LetDec rep)
    pat StmAux (ExpDec rep)
    aux Exp rep
    e) [Stm rep] -> CSEM rep a
    m = do
      CSEState (ExpressionSubstitutions rep
    esubsts, NameSubstitutions
    nsubsts) Bool
    cse_arrays <- ReaderT (CSEState rep) Identity (CSEState rep)
    forall r (m :: * -> *). MonadReader r m => m r
    ask
      let e' :: Exp rep
    e' = NameSubstitutions -> Exp rep -> Exp rep
    forall a. Substitute a => NameSubstitutions -> a -> a
    substituteNames NameSubstitutions
    nsubsts Exp rep
    e
          pat' :: Pat (LetDec rep)
    pat' = NameSubstitutions -> Pat (LetDec rep) -> Pat (LetDec rep)
    forall a. Substitute a => NameSubstitutions -> a -> a
    substituteNames NameSubstitutions
    nsubsts Pat (LetDec rep)
    pat
      if Bool -> Bool
    not (Exp rep -> Bool
    forall {rep}. Exp rep -> Bool
    alreadyAliases Exp rep
    e) Bool -> Bool -> Bool
    && (PatElem (LetDec rep) -> Bool) -> [PatElem (LetDec rep)] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any (Bool -> PatElem (LetDec rep) -> Bool
    forall {dec}. Typed dec => Bool -> PatElem dec -> Bool
    bad Bool
    cse_arrays) (Pat (LetDec rep) -> [PatElem (LetDec rep)]
    forall dec. Pat dec -> [PatElem dec]
    patElems Pat (LetDec rep)
    pat)
        then [Stm rep] -> CSEM rep a
    m [Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
    forall rep.
    Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
    Let Pat (LetDec rep)
    pat' StmAux (ExpDec rep)
    aux Exp rep
    e']
        else case (ExpDec rep, Exp rep)
    -> ExpressionSubstitutions rep -> Maybe (Certs, Pat (LetDec rep))
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup (StmAux (ExpDec rep) -> ExpDec rep
    forall dec. StmAux dec -> dec
    stmAuxDec StmAux (ExpDec rep)
    aux, Exp rep
    e') ExpressionSubstitutions rep
    esubsts of
          Just (Certs
    subcs, Pat (LetDec rep)
    subpat) -> do
            let subsumes :: Bool
    subsumes = (VName -> Bool) -> [VName] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all (VName -> [VName] -> Bool
    forall a. Eq a => a -> [a] -> Bool
    forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
    `elem` Certs -> [VName]
    unCerts Certs
    subcs) (Certs -> [VName]
    unCerts (StmAux (ExpDec rep) -> Certs
    forall dec. StmAux dec -> Certs
    stmAuxCerts StmAux (ExpDec rep)
    aux))
            -- We can only do a plain name substitution if it doesn't
            -- violate any certificate dependencies.
            (CSEState rep -> CSEState rep) -> CSEM rep a -> CSEM rep a
    forall a.
    (CSEState rep -> CSEState rep)
    -> ReaderT (CSEState rep) Identity a
    -> ReaderT (CSEState rep) Identity a
    forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
    local (if Bool
    subsumes then Pat (LetDec rep)
    -> Pat (LetDec rep) -> CSEState rep -> CSEState rep
    forall dec rep. Pat dec -> Pat dec -> CSEState rep -> CSEState rep
    addNameSubst Pat (LetDec rep)
    pat' Pat (LetDec rep)
    subpat else CSEState rep -> CSEState rep
    forall a. a -> a
    id) (CSEM rep a -> CSEM rep a) -> CSEM rep a -> CSEM rep a
    forall a b. (a -> b) -> a -> b
    $ do
              let lets :: [Stm rep]
    lets =
                    [ Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
    forall rep.
    Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
    Let ([PatElem (LetDec rep)] -> Pat (LetDec rep)
    forall dec. [PatElem dec] -> Pat dec
    Pat [PatElem (LetDec rep)
    patElem']) StmAux (ExpDec rep)
    aux (Exp rep -> Stm rep) -> Exp rep -> Stm rep
    forall a b. (a -> b) -> a -> b
    $
                        BasicOp -> Exp rep
    forall rep. BasicOp -> Exp rep
    BasicOp (SubExp -> BasicOp
    SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
    forall a b. (a -> b) -> a -> b
    $ VName -> SubExp
    Var (VName -> SubExp) -> VName -> SubExp
    forall a b. (a -> b) -> a -> b
    $ PatElem (LetDec rep) -> VName
    forall dec. PatElem dec -> VName
    patElemName PatElem (LetDec rep)
    patElem)
                      | (VName
    name, PatElem (LetDec rep)
    patElem) <- [VName]
    -> [PatElem (LetDec rep)] -> [(VName, PatElem (LetDec rep))]
    forall a b. [a] -> [b] -> [(a, b)]
    zip (Pat (LetDec rep) -> [VName]
    forall dec. Pat dec -> [VName]
    patNames Pat (LetDec rep)
    pat') ([PatElem (LetDec rep)] -> [(VName, PatElem (LetDec rep))])
    -> [PatElem (LetDec rep)] -> [(VName, PatElem (LetDec rep))]
    forall a b. (a -> b) -> a -> b
    $ Pat (LetDec rep) -> [PatElem (LetDec rep)]
    forall dec. Pat dec -> [PatElem dec]
    patElems Pat (LetDec rep)
    subpat,
                        let patElem' :: PatElem (LetDec rep)
    patElem' = PatElem (LetDec rep)
    patElem {patElemName = name}
                    ]
              [Stm rep] -> CSEM rep a
    m [Stm rep]
    lets
          Maybe (Certs, Pat (LetDec rep))
    _ ->
            (CSEState rep -> CSEState rep) -> CSEM rep a -> CSEM rep a
    forall a.
    (CSEState rep -> CSEState rep)
    -> ReaderT (CSEState rep) Identity a
    -> ReaderT (CSEState rep) Identity a
    forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
    local (Pat (LetDec rep)
    -> ExpDec rep -> Certs -> Exp rep -> CSEState rep -> CSEState rep
    forall rep.
    ASTRep rep =>
    Pat (LetDec rep)
    -> ExpDec rep -> Certs -> Exp rep -> CSEState rep -> CSEState rep
    addExpSubst Pat (LetDec rep)
    pat' (StmAux (ExpDec rep) -> ExpDec rep
    forall dec. StmAux dec -> dec
    stmAuxDec StmAux (ExpDec rep)
    aux) (StmAux (ExpDec rep) -> Certs
    forall dec. StmAux dec -> Certs
    stmAuxCerts StmAux (ExpDec rep)
    aux) Exp rep
    e') (CSEM rep a -> CSEM rep a) -> CSEM rep a -> CSEM rep a
    forall a b. (a -> b) -> a -> b
    $
              [Stm rep] -> CSEM rep a
    m [Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
    forall rep.
    Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
    Let Pat (LetDec rep)
    pat' StmAux (ExpDec rep)
    aux Exp rep
    e']
      where
        alreadyAliases :: Exp rep -> Bool
    alreadyAliases (BasicOp Index {}) = Bool
    True
        alreadyAliases (BasicOp Reshape {}) = Bool
    True
        alreadyAliases Exp rep
    _ = Bool
    False
        bad :: Bool -> PatElem dec -> Bool
    bad Bool
    cse_arrays PatElem dec
    pe
          | Mem {} <- PatElem dec -> Type
    forall dec. Typed dec => PatElem dec -> Type
    patElemType PatElem dec
    pe = Bool
    True
          | Array {} <- PatElem dec -> Type
    forall dec. Typed dec => PatElem dec -> Type
    patElemType PatElem dec
    pe, Bool -> Bool
    not Bool
    cse_arrays = Bool
    True
          | PatElem dec -> VName
    forall dec. PatElem dec -> VName
    patElemName PatElem dec
    pe VName -> Names -> Bool
    `nameIn` Names
    consumed = Bool
    True
          | Bool
    otherwise = Bool
    False

    type ExpressionSubstitutions rep =
      M.Map
        (ExpDec rep, Exp rep)
        (Certs, Pat (LetDec rep))

    type NameSubstitutions = M.Map VName VName

    data CSEState rep = CSEState
      { forall rep.
    CSEState rep -> (ExpressionSubstitutions rep, NameSubstitutions)
    _cseSubstitutions :: (ExpressionSubstitutions rep, NameSubstitutions),
        forall rep. CSEState rep -> Bool
    _cseArrays :: Bool
      }

    newCSEState :: Bool -> CSEState rep
    newCSEState :: forall rep. Bool -> CSEState rep
    newCSEState = (ExpressionSubstitutions rep, NameSubstitutions)
    -> Bool -> CSEState rep
    forall rep.
    (ExpressionSubstitutions rep, NameSubstitutions)
    -> Bool -> CSEState rep
    CSEState (ExpressionSubstitutions rep
    forall k a. Map k a
    M.empty, NameSubstitutions
    forall k a. Map k a
    M.empty)

    mkSubsts :: Pat dec -> Pat dec -> M.Map VName VName
    mkSubsts :: forall dec. Pat dec -> Pat dec -> NameSubstitutions
    mkSubsts Pat dec
    pat Pat dec
    vs = [(VName, VName)] -> NameSubstitutions
    forall k a. Ord k => [(k, a)] -> Map k a
    M.fromList ([(VName, VName)] -> NameSubstitutions)
    -> [(VName, VName)] -> NameSubstitutions
    forall a b. (a -> b) -> a -> b
    $ [VName] -> [VName] -> [(VName, VName)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip (Pat dec -> [VName]
    forall dec. Pat dec -> [VName]
    patNames Pat dec
    pat) (Pat dec -> [VName]
    forall dec. Pat dec -> [VName]
    patNames Pat dec
    vs)

    addNameSubst :: Pat dec -> Pat dec -> CSEState rep -> CSEState rep
    addNameSubst :: forall dec rep. Pat dec -> Pat dec -> CSEState rep -> CSEState rep
    addNameSubst Pat dec
    pat Pat dec
    subpat (CSEState (ExpressionSubstitutions rep
    esubsts, NameSubstitutions
    nsubsts) Bool
    cse_arrays) =
      (ExpressionSubstitutions rep, NameSubstitutions)
    -> Bool -> CSEState rep
    forall rep.
    (ExpressionSubstitutions rep, NameSubstitutions)
    -> Bool -> CSEState rep
    CSEState (ExpressionSubstitutions rep
    esubsts, Pat dec -> Pat dec -> NameSubstitutions
    forall dec. Pat dec -> Pat dec -> NameSubstitutions
    mkSubsts Pat dec
    pat Pat dec
    subpat NameSubstitutions -> NameSubstitutions -> NameSubstitutions
    forall k a. Ord k => Map k a -> Map k a -> Map k a
    `M.union` NameSubstitutions
    nsubsts) Bool
    cse_arrays

    addExpSubst ::
      (ASTRep rep) =>
      Pat (LetDec rep) ->
      ExpDec rep ->
      Certs ->
      Exp rep ->
      CSEState rep ->
      CSEState rep
    addExpSubst :: forall rep.
    ASTRep rep =>
    Pat (LetDec rep)
    -> ExpDec rep -> Certs -> Exp rep -> CSEState rep -> CSEState rep
    addExpSubst Pat (LetDec rep)
    pat ExpDec rep
    edec Certs
    cs Exp rep
    e (CSEState (ExpressionSubstitutions rep
    esubsts, NameSubstitutions
    nsubsts) Bool
    cse_arrays) =
      (ExpressionSubstitutions rep, NameSubstitutions)
    -> Bool -> CSEState rep
    forall rep.
    (ExpressionSubstitutions rep, NameSubstitutions)
    -> Bool -> CSEState rep
    CSEState ((ExpDec rep, Exp rep)
    -> (Certs, Pat (LetDec rep))
    -> ExpressionSubstitutions rep
    -> ExpressionSubstitutions rep
    forall k a. Ord k => k -> a -> Map k a -> Map k a
    M.insert (ExpDec rep
    edec, Exp rep
    e) (Certs
    cs, Pat (LetDec rep)
    pat) ExpressionSubstitutions rep
    esubsts, NameSubstitutions
    nsubsts) Bool
    cse_arrays

    -- | The operations that permit CSE.
    class CSEInOp op where
      -- | Perform CSE within any nested expressions.
      cseInOp :: op -> CSEM rep op

    instance CSEInOp (NoOp rep) where
      cseInOp :: forall rep. NoOp rep -> CSEM rep (NoOp rep)
    cseInOp NoOp rep
    NoOp = NoOp rep -> ReaderT (CSEState rep) Identity (NoOp rep)
    forall a. a -> ReaderT (CSEState rep) Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure NoOp rep
    forall {k} (rep :: k). NoOp rep
    NoOp

    subCSE :: CSEM rep r -> CSEM otherrep r
    subCSE :: forall rep r otherrep. CSEM rep r -> CSEM otherrep r
    subCSE CSEM rep r
    m = do
      CSEState (ExpressionSubstitutions otherrep, NameSubstitutions)
    _ Bool
    cse_arrays <- ReaderT (CSEState otherrep) Identity (CSEState otherrep)
    forall r (m :: * -> *). MonadReader r m => m r
    ask
      r -> CSEM otherrep r
    forall a. a -> ReaderT (CSEState otherrep) Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (r -> CSEM otherrep r) -> r -> CSEM otherrep r
    forall a b. (a -> b) -> a -> b
    $ CSEM rep r -> CSEState rep -> r
    forall r a. Reader r a -> r -> a
    runReader CSEM rep r
    m (CSEState rep -> r) -> CSEState rep -> r
    forall a b. (a -> b) -> a -> b
    $ Bool -> CSEState rep
    forall rep. Bool -> CSEState rep
    newCSEState Bool
    cse_arrays

    instance
      ( Aliased rep,
        CSEInOp (Op rep),
        CSEInOp (op rep)
      ) =>
      CSEInOp (GPU.HostOp op rep)
      where
      cseInOp :: forall rep. HostOp op rep -> CSEM rep (HostOp op rep)
    cseInOp (GPU.SegOp SegOp SegLevel rep
    op) = SegOp SegLevel rep -> HostOp op rep
    forall (op :: * -> *) rep. SegOp SegLevel rep -> HostOp op rep
    GPU.SegOp (SegOp SegLevel rep -> HostOp op rep)
    -> ReaderT (CSEState rep) Identity (SegOp SegLevel rep)
    -> ReaderT (CSEState rep) Identity (HostOp op rep)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> SegOp SegLevel rep
    -> ReaderT (CSEState rep) Identity (SegOp SegLevel rep)
    forall rep. SegOp SegLevel rep -> CSEM rep (SegOp SegLevel rep)
    forall op rep. CSEInOp op => op -> CSEM rep op
    cseInOp SegOp SegLevel rep
    op
      cseInOp (GPU.OtherOp op rep
    op) = op rep -> HostOp op rep
    forall (op :: * -> *) rep. op rep -> HostOp op rep
    GPU.OtherOp (op rep -> HostOp op rep)
    -> ReaderT (CSEState rep) Identity (op rep)
    -> ReaderT (CSEState rep) Identity (HostOp op rep)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> op rep -> ReaderT (CSEState rep) Identity (op rep)
    forall rep. op rep -> CSEM rep (op rep)
    forall op rep. CSEInOp op => op -> CSEM rep op
    cseInOp op rep
    op
      cseInOp (GPU.GPUBody [Type]
    types Body rep
    body) =
        CSEM rep (HostOp op rep)
    -> ReaderT (CSEState rep) Identity (HostOp op rep)
    forall rep r otherrep. CSEM rep r -> CSEM otherrep r
    subCSE (CSEM rep (HostOp op rep)
     -> ReaderT (CSEState rep) Identity (HostOp op rep))
    -> CSEM rep (HostOp op rep)
    -> ReaderT (CSEState rep) Identity (HostOp op rep)
    forall a b. (a -> b) -> a -> b
    $ [Type] -> Body rep -> HostOp op rep
    forall (op :: * -> *) rep. [Type] -> Body rep -> HostOp op rep
    GPU.GPUBody [Type]
    types (Body rep -> HostOp op rep)
    -> ReaderT (CSEState rep) Identity (Body rep)
    -> CSEM rep (HostOp op rep)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> [Diet] -> Body rep -> ReaderT (CSEState rep) Identity (Body rep)
    forall rep.
    (Aliased rep, CSEInOp (Op rep)) =>
    [Diet] -> Body rep -> CSEM rep (Body rep)
    cseInBody ((Type -> Diet) -> [Type] -> [Diet]
    forall a b. (a -> b) -> [a] -> [b]
    map (Diet -> Type -> Diet
    forall a b. a -> b -> a
    const Diet
    Observe) [Type]
    types) Body rep
    body
      cseInOp HostOp op rep
    x = HostOp op rep -> ReaderT (CSEState rep) Identity (HostOp op rep)
    forall a. a -> ReaderT (CSEState rep) Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure HostOp op rep
    x

    instance
      ( Aliased rep,
        CSEInOp (Op rep),
        CSEInOp (op rep)
      ) =>
      CSEInOp (MC.MCOp op rep)
      where
      cseInOp :: forall rep. MCOp op rep -> CSEM rep (MCOp op rep)
    cseInOp (MC.ParOp Maybe (SegOp () rep)
    par_op SegOp () rep
    op) =
        Maybe (SegOp () rep) -> SegOp () rep -> MCOp op rep
    forall (op :: * -> *) rep.
    Maybe (SegOp () rep) -> SegOp () rep -> MCOp op rep
    MC.ParOp (Maybe (SegOp () rep) -> SegOp () rep -> MCOp op rep)
    -> ReaderT (CSEState rep) Identity (Maybe (SegOp () rep))
    -> ReaderT (CSEState rep) Identity (SegOp () rep -> MCOp op rep)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (SegOp () rep -> ReaderT (CSEState rep) Identity (SegOp () rep))
    -> Maybe (SegOp () rep)
    -> ReaderT (CSEState rep) Identity (Maybe (SegOp () rep))
    forall (t :: * -> *) (f :: * -> *) a b.
    (Traversable t, Applicative f) =>
    (a -> f b) -> t a -> f (t b)
    forall (f :: * -> *) a b.
    Applicative f =>
    (a -> f b) -> Maybe a -> f (Maybe b)
    traverse SegOp () rep -> ReaderT (CSEState rep) Identity (SegOp () rep)
    forall rep. SegOp () rep -> CSEM rep (SegOp () rep)
    forall op rep. CSEInOp op => op -> CSEM rep op
    cseInOp Maybe (SegOp () rep)
    par_op ReaderT (CSEState rep) Identity (SegOp () rep -> MCOp op rep)
    -> ReaderT (CSEState rep) Identity (SegOp () rep)
    -> ReaderT (CSEState rep) Identity (MCOp op rep)
    forall a b.
    ReaderT (CSEState rep) Identity (a -> b)
    -> ReaderT (CSEState rep) Identity a
    -> ReaderT (CSEState rep) Identity b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> SegOp () rep -> ReaderT (CSEState rep) Identity (SegOp () rep)
    forall rep. SegOp () rep -> CSEM rep (SegOp () rep)
    forall op rep. CSEInOp op => op -> CSEM rep op
    cseInOp SegOp () rep
    op
      cseInOp (MC.OtherOp op rep
    op) =
        op rep -> MCOp op rep
    forall (op :: * -> *) rep. op rep -> MCOp op rep
    MC.OtherOp (op rep -> MCOp op rep)
    -> ReaderT (CSEState rep) Identity (op rep)
    -> ReaderT (CSEState rep) Identity (MCOp op rep)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> op rep -> ReaderT (CSEState rep) Identity (op rep)
    forall rep. op rep -> CSEM rep (op rep)
    forall op rep. CSEInOp op => op -> CSEM rep op
    cseInOp op rep
    op

    instance
      (Aliased rep, CSEInOp (Op rep)) =>
      CSEInOp (GPU.SegOp lvl rep)
      where
      cseInOp :: forall rep. SegOp lvl rep -> CSEM rep (SegOp lvl rep)
    cseInOp =
        CSEM rep (SegOp lvl rep) -> CSEM rep (SegOp lvl rep)
    forall rep r otherrep. CSEM rep r -> CSEM otherrep r
    subCSE
          (CSEM rep (SegOp lvl rep) -> CSEM rep (SegOp lvl rep))
    -> (SegOp lvl rep -> CSEM rep (SegOp lvl rep))
    -> SegOp lvl rep
    -> CSEM rep (SegOp lvl rep)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SegOpMapper lvl rep rep (ReaderT (CSEState rep) Identity)
    -> SegOp lvl rep -> CSEM rep (SegOp lvl rep)
    forall (m :: * -> *) lvl frep trep.
    Monad m =>
    SegOpMapper lvl frep trep m -> SegOp lvl frep -> m (SegOp lvl trep)
    GPU.mapSegOpM
            ((SubExp -> ReaderT (CSEState rep) Identity SubExp)
    -> (Lambda rep -> ReaderT (CSEState rep) Identity (Lambda rep))
    -> (KernelBody rep
        -> ReaderT (CSEState rep) Identity (KernelBody rep))
    -> (VName -> ReaderT (CSEState rep) Identity VName)
    -> (lvl -> ReaderT (CSEState rep) Identity lvl)
    -> SegOpMapper lvl rep rep (ReaderT (CSEState rep) Identity)
    forall lvl frep trep (m :: * -> *).
    (SubExp -> m SubExp)
    -> (Lambda frep -> m (Lambda trep))
    -> (KernelBody frep -> m (KernelBody trep))
    -> (VName -> m VName)
    -> (lvl -> m lvl)
    -> SegOpMapper lvl frep trep m
    GPU.SegOpMapper SubExp -> ReaderT (CSEState rep) Identity SubExp
    forall a. a -> ReaderT (CSEState rep) Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Lambda rep -> ReaderT (CSEState rep) Identity (Lambda rep)
    forall rep.
    (Aliased rep, CSEInOp (Op rep)) =>
    Lambda rep -> CSEM rep (Lambda rep)
    cseInLambda KernelBody rep -> ReaderT (CSEState rep) Identity (KernelBody rep)
    forall rep.
    (Aliased rep, CSEInOp (Op rep)) =>
    KernelBody rep -> CSEM rep (KernelBody rep)
    cseInKernelBody VName -> ReaderT (CSEState rep) Identity VName
    forall a. a -> ReaderT (CSEState rep) Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure lvl -> ReaderT (CSEState rep) Identity lvl
    forall a. a -> ReaderT (CSEState rep) Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure)

    cseInKernelBody ::
      (Aliased rep, CSEInOp (Op rep)) =>
      GPU.KernelBody rep ->
      CSEM rep (GPU.KernelBody rep)
    cseInKernelBody :: forall rep.
    (Aliased rep, CSEInOp (Op rep)) =>
    KernelBody rep -> CSEM rep (KernelBody rep)
    cseInKernelBody (GPU.KernelBody BodyDec rep
    bodydec Stms rep
    stms [KernelResult]
    res) = do
      Body BodyDec rep
    _ Stms rep
    stms' Result
    _ <- [Diet] -> Body rep -> CSEM rep (Body rep)
    forall rep.
    (Aliased rep, CSEInOp (Op rep)) =>
    [Diet] -> Body rep -> CSEM rep (Body rep)
    cseInBody ((KernelResult -> Diet) -> [KernelResult] -> [Diet]
    forall a b. (a -> b) -> [a] -> [b]
    map (Diet -> KernelResult -> Diet
    forall a b. a -> b -> a
    const Diet
    Observe) [KernelResult]
    res) (Body rep -> CSEM rep (Body rep))
    -> Body rep -> CSEM rep (Body rep)
    forall a b. (a -> b) -> a -> b
    $ BodyDec rep -> Stms rep -> Result -> Body rep
    forall rep. BodyDec rep -> Stms rep -> Result -> Body rep
    Body BodyDec rep
    bodydec Stms rep
    stms []
      KernelBody rep -> CSEM rep (KernelBody rep)
    forall a. a -> ReaderT (CSEState rep) Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (KernelBody rep -> CSEM rep (KernelBody rep))
    -> KernelBody rep -> CSEM rep (KernelBody rep)
    forall a b. (a -> b) -> a -> b
    $ BodyDec rep -> Stms rep -> [KernelResult] -> KernelBody rep
    forall rep.
    BodyDec rep -> Stms rep -> [KernelResult] -> KernelBody rep
    GPU.KernelBody BodyDec rep
    bodydec Stms rep
    stms' [KernelResult]
    res

    instance (CSEInOp (op rep)) => CSEInOp (Memory.MemOp op rep) where
      cseInOp :: forall rep. MemOp op rep -> CSEM rep (MemOp op rep)
    cseInOp o :: MemOp op rep
    o@Memory.Alloc {} = MemOp op rep -> ReaderT (CSEState rep) Identity (MemOp op rep)
    forall a. a -> ReaderT (CSEState rep) Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure MemOp op rep
    o
      cseInOp (Memory.Inner op rep
    k) = op rep -> MemOp op rep
    forall (inner :: * -> *) rep. inner rep -> MemOp inner rep
    Memory.Inner (op rep -> MemOp op rep)
    -> ReaderT (CSEState rep) Identity (op rep)
    -> ReaderT (CSEState rep) Identity (MemOp op rep)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> CSEM Any (op rep) -> ReaderT (CSEState rep) Identity (op rep)
    forall rep r otherrep. CSEM rep r -> CSEM otherrep r
    subCSE (op rep -> CSEM Any (op rep)
    forall rep. op rep -> CSEM rep (op rep)
    forall op rep. CSEInOp op => op -> CSEM rep op
    cseInOp op rep
    k)

    instance
      (AliasableRep rep, CSEInOp (Op (Aliases rep))) =>
      CSEInOp (SOAC.SOAC (Aliases rep))
      where
      cseInOp :: forall rep. SOAC (Aliases rep) -> CSEM rep (SOAC (Aliases rep))
    cseInOp = CSEM (Aliases rep) (SOAC (Aliases rep))
    -> CSEM rep (SOAC (Aliases rep))
    forall rep r otherrep. CSEM rep r -> CSEM otherrep r
    subCSE (CSEM (Aliases rep) (SOAC (Aliases rep))
     -> CSEM rep (SOAC (Aliases rep)))
    -> (SOAC (Aliases rep) -> CSEM (Aliases rep) (SOAC (Aliases rep)))
    -> SOAC (Aliases rep)
    -> CSEM rep (SOAC (Aliases rep))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . SOACMapper
      (Aliases rep)
      (Aliases rep)
      (ReaderT (CSEState (Aliases rep)) Identity)
    -> SOAC (Aliases rep) -> CSEM (Aliases rep) (SOAC (Aliases rep))
    forall (m :: * -> *) frep trep.
    Monad m =>
    SOACMapper frep trep m -> SOAC frep -> m (SOAC trep)
    SOAC.mapSOACM ((SubExp -> ReaderT (CSEState (Aliases rep)) Identity SubExp)
    -> (Lambda (Aliases rep)
        -> ReaderT
             (CSEState (Aliases rep)) Identity (Lambda (Aliases rep)))
    -> (VName -> ReaderT (CSEState (Aliases rep)) Identity VName)
    -> SOACMapper
         (Aliases rep)
         (Aliases rep)
         (ReaderT (CSEState (Aliases rep)) Identity)
    forall frep trep (m :: * -> *).
    (SubExp -> m SubExp)
    -> (Lambda frep -> m (Lambda trep))
    -> (VName -> m VName)
    -> SOACMapper frep trep m
    SOAC.SOACMapper SubExp -> ReaderT (CSEState (Aliases rep)) Identity SubExp
    forall a. a -> ReaderT (CSEState (Aliases rep)) Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Lambda (Aliases rep)
    -> ReaderT (CSEState (Aliases rep)) Identity (Lambda (Aliases rep))
    forall rep.
    (Aliased rep, CSEInOp (Op rep)) =>
    Lambda rep -> CSEM rep (Lambda rep)
    cseInLambda VName -> ReaderT (CSEState (Aliases rep)) Identity VName
    forall a. a -> ReaderT (CSEState (Aliases rep)) Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure)