    {-# LANGUAGE TypeFamilies #-}

    module Futhark.AD.Rev.SOAC (vjpSOAC) where

    import Control.Monad
    import Futhark.AD.Rev.Hist
    import Futhark.AD.Rev.Map
    import Futhark.AD.Rev.Monad
    import Futhark.AD.Rev.Reduce
    import Futhark.AD.Rev.Scan
    import Futhark.AD.Rev.Scatter
    import Futhark.Analysis.PrimExp.Convert
    import Futhark.Builder
    import Futhark.IR.SOACS
    import Futhark.Tools
    import Futhark.Util (chunks)

    -- We split any multi-op scan or reduction into multiple operations so
    -- we can detect special cases.  Post-AD, the result may be fused
    -- again.
    splitScanRed ::
      VjpOps ->
      ([a] -> ADM (ScremaForm SOACS), a -> [SubExp]) ->
      (Pat Type, StmAux (), [a], SubExp, [VName]) ->
      ADM () ->
      ADM ()
    splitScanRed :: forall a.
    VjpOps
    -> ([a] -> ADM (ScremaForm SOACS), a -> [SubExp])
    -> (Pat Type, StmAux (), [a], SubExp, [VName])
    -> ADM ()
    -> ADM ()
    splitScanRed VjpOps
    vjpops ([a] -> ADM (ScremaForm SOACS)
    opSOAC, a -> [SubExp]
    opNeutral) (Pat Type
    pat, StmAux ()
    aux, [a]
    ops, SubExp
    w, [VName]
    as) ADM ()
    m = do
      let ks :: [Int]
    ks = (a -> Int) -> [a] -> [Int]
    forall a b. (a -> b) -> [a] -> [b]
    map ([SubExp] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length ([SubExp] -> Int) -> (a -> [SubExp]) -> a -> Int
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . a -> [SubExp]
    opNeutral) [a]
    ops
          pat_per_op :: [Pat Type]
    pat_per_op = ([PatElem Type] -> Pat Type) -> [[PatElem Type]] -> [Pat Type]
    forall a b. (a -> b) -> [a] -> [b]
    map [PatElem Type] -> Pat Type
    forall dec. [PatElem dec] -> Pat dec
    Pat ([[PatElem Type]] -> [Pat Type]) -> [[PatElem Type]] -> [Pat Type]
    forall a b. (a -> b) -> a -> b
    $ [Int] -> [PatElem Type] -> [[PatElem Type]]
    forall a. [Int] -> [a] -> [[a]]
    chunks [Int]
    ks ([PatElem Type] -> [[PatElem Type]])
    -> [PatElem Type] -> [[PatElem Type]]
    forall a b. (a -> b) -> a -> b
    $ Pat Type -> [PatElem Type]
    forall dec. Pat dec -> [PatElem dec]
    patElems Pat Type
    pat
          as_per_op :: [[VName]]
    as_per_op = [Int] -> [VName] -> [[VName]]
    forall a. [Int] -> [a] -> [[a]]
    chunks [Int]
    ks [VName]
    as
          onOps :: [a] -> [Pat Type] -> [[VName]] -> ADM ()
    onOps (a
    op : [a]
    ops') (Pat Type
    op_pat : [Pat Type]
    op_pats') ([VName]
    op_as : [[VName]]
    op_as') = do
            op_form <- [a] -> ADM (ScremaForm SOACS)
    opSOAC [a
    op]
            vjpSOAC vjpops op_pat aux (Screma w op_as op_form) $
              onOps ops' op_pats' op_as'
          onOps [a]
    _ [Pat Type]
    _ [[VName]]
    _ = ADM ()
    m
      [a] -> [Pat Type] -> [[VName]] -> ADM ()
    onOps [a]
    ops [Pat Type]
    pat_per_op [[VName]]
    as_per_op

    -- We split multi-op histograms into multiple operations so we
    -- can take advantage of special cases. Post-AD, the result may
    -- be fused again.
    splitHist :: VjpOps -> Pat Type -> StmAux () -> [HistOp SOACS] -> SubExp -> [VName] -> [VName] -> ADM () -> ADM ()
    splitHist :: VjpOps
    -> Pat Type
    -> StmAux ()
    -> [HistOp SOACS]
    -> SubExp
    -> [VName]
    -> [VName]
    -> ADM ()
    -> ADM ()
    splitHist VjpOps
    vjpops Pat Type
    pat StmAux ()
    aux [HistOp SOACS]
    ops SubExp
    w [VName]
    is [VName]
    as ADM ()
    m = do
      let ks :: [Int]
    ks = (HistOp SOACS -> Int) -> [HistOp SOACS] -> [Int]
    forall a b. (a -> b) -> [a] -> [b]
    map ([SubExp] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length ([SubExp] -> Int)
    -> (HistOp SOACS -> [SubExp]) -> HistOp SOACS -> Int
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . HistOp SOACS -> [SubExp]
    forall rep. HistOp rep -> [SubExp]
    histNeutral) [HistOp SOACS]
    ops
          pat_per_op :: [Pat Type]
    pat_per_op = ([PatElem Type] -> Pat Type) -> [[PatElem Type]] -> [Pat Type]
    forall a b. (a -> b) -> [a] -> [b]
    map [PatElem Type] -> Pat Type
    forall dec. [PatElem dec] -> Pat dec
    Pat ([[PatElem Type]] -> [Pat Type]) -> [[PatElem Type]] -> [Pat Type]
    forall a b. (a -> b) -> a -> b
    $ [Int] -> [PatElem Type] -> [[PatElem Type]]
    forall a. [Int] -> [a] -> [[a]]
    chunks [Int]
    ks ([PatElem Type] -> [[PatElem Type]])
    -> [PatElem Type] -> [[PatElem Type]]
    forall a b. (a -> b) -> a -> b
    $ Pat Type -> [PatElem Type]
    forall dec. Pat dec -> [PatElem dec]
    patElems Pat Type
    pat
          as_per_op :: [[VName]]
    as_per_op = [Int] -> [VName] -> [[VName]]
    forall a. [Int] -> [a] -> [[a]]
    chunks [Int]
    ks [VName]
    as
          onOps :: [HistOp SOACS] -> [Pat Type] -> [VName] -> [[VName]] -> ADM ()
    onOps (HistOp SOACS
    op : [HistOp SOACS]
    ops') (Pat Type
    op_pat : [Pat Type]
    op_pats') (VName
    op_is : [VName]
    op_is') ([VName]
    op_as : [[VName]]
    op_as') = do
            f <- [Type] -> ADM (Lambda SOACS)
    forall rep (m :: * -> *).
    (Buildable rep, MonadFreshNames m) =>
    [Type] -> m (Lambda rep)
    mkIdentityLambda ([Type] -> ADM (Lambda SOACS))
    -> ([Type] -> [Type]) -> [Type] -> ADM (Lambda SOACS)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (PrimType -> Type
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    int64 Type -> [Type] -> [Type]
    forall a. a -> [a] -> [a]
    :) ([Type] -> ADM (Lambda SOACS)) -> ADM [Type] -> ADM (Lambda SOACS)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< (VName -> ADM Type) -> [VName] -> ADM [Type]
    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) -> [a] -> f [b]
    traverse VName -> ADM Type
    forall rep (m :: * -> *). HasScope rep m => VName -> m Type
    lookupType [VName]
    op_as
            vjpSOAC vjpops op_pat aux (Hist w (op_is : op_as) [op] f) $
              onOps ops' op_pats' op_is' op_as'
          onOps [HistOp SOACS]
    _ [Pat Type]
    _ [VName]
    _ [[VName]]
    _ = ADM ()
    m
      [HistOp SOACS] -> [Pat Type] -> [VName] -> [[VName]] -> ADM ()
    onOps [HistOp SOACS]
    ops [Pat Type]
    pat_per_op [VName]
    is [[VName]]
    as_per_op

    -- unfusing a map-histogram construct into a map and a histogram.
    histomapToMapAndHist :: Pat Type -> (SubExp, [HistOp SOACS], Lambda SOACS, [VName]) -> ADM (Stm SOACS, Stm SOACS)
    histomapToMapAndHist :: Pat Type
    -> (SubExp, [HistOp SOACS], Lambda SOACS, [VName])
    -> ADM (Stm SOACS, Stm SOACS)
    histomapToMapAndHist (Pat [PatElem Type]
    pes) (SubExp
    w, [HistOp SOACS]
    histops, Lambda SOACS
    map_lam, [VName]
    as) = do
      map_pat <- (Type -> ADM Ident) -> [Type] -> ADM [Ident]
    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) -> [a] -> f [b]
    traverse Type -> ADM Ident
    accMapPatElem ([Type] -> ADM [Ident]) -> [Type] -> ADM [Ident]
    forall a b. (a -> b) -> a -> b
    $ Lambda SOACS -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda SOACS
    map_lam
      let map_stm = [Ident] -> Exp SOACS -> Stm SOACS
    forall rep. Buildable rep => [Ident] -> Exp rep -> Stm rep
    mkLet [Ident]
    map_pat (Exp SOACS -> Stm SOACS) -> Exp SOACS -> Stm SOACS
    forall a b. (a -> b) -> a -> b
    $ Op SOACS -> Exp SOACS
    forall rep. Op rep -> Exp rep
    Op (Op SOACS -> Exp SOACS) -> Op SOACS -> Exp SOACS
    forall a b. (a -> b) -> a -> b
    $ SubExp -> [VName] -> ScremaForm SOACS -> SOAC SOACS
    forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
    Screma SubExp
    w [VName]
    as (ScremaForm SOACS -> SOAC SOACS) -> ScremaForm SOACS -> SOAC SOACS
    forall a b. (a -> b) -> a -> b
    $ Lambda SOACS -> ScremaForm SOACS
    forall rep. Lambda rep -> ScremaForm rep
    mapSOAC Lambda SOACS
    map_lam
      new_lam <- mkIdentityLambda $ lambdaReturnType map_lam
      let hist_stm = Pat (LetDec SOACS)
    -> StmAux (ExpDec SOACS) -> Exp SOACS -> Stm SOACS
    forall rep.
    Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
    Let ([PatElem Type] -> Pat Type
    forall dec. [PatElem dec] -> Pat dec
    Pat [PatElem Type]
    pes) (() -> StmAux ()
    forall dec. dec -> StmAux dec
    defAux ()) (Exp SOACS -> Stm SOACS) -> Exp SOACS -> Stm SOACS
    forall a b. (a -> b) -> a -> b
    $ Op SOACS -> Exp SOACS
    forall rep. Op rep -> Exp rep
    Op (Op SOACS -> Exp SOACS) -> Op SOACS -> Exp SOACS
    forall a b. (a -> b) -> a -> b
    $ SubExp -> [VName] -> [HistOp SOACS] -> Lambda SOACS -> SOAC SOACS
    forall rep.
    SubExp -> [VName] -> [HistOp rep] -> Lambda rep -> SOAC rep
    Hist SubExp
    w ((Ident -> VName) -> [Ident] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map Ident -> VName
    identName [Ident]
    map_pat) [HistOp SOACS]
    histops Lambda SOACS
    new_lam
      pure (map_stm, hist_stm)
      where
        accMapPatElem :: Type -> ADM Ident
    accMapPatElem =
          [Char] -> Type -> ADM Ident
    forall (m :: * -> *).
    MonadFreshNames m =>
    [Char] -> Type -> m Ident
    newIdent [Char]
    "hist_map_res" (Type -> ADM Ident) -> (Type -> Type) -> Type -> ADM Ident
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Type -> SubExp -> Type
    forall d.
    ArrayShape (ShapeBase d) =>
    TypeBase (ShapeBase d) NoUniqueness
    -> d -> TypeBase (ShapeBase d) NoUniqueness
    `arrayOfRow` SubExp
    w)

    commonSOAC :: Pat Type -> StmAux () -> SOAC SOACS -> ADM () -> ADM [Adj]
    commonSOAC :: Pat Type -> StmAux () -> SOAC SOACS -> ADM () -> ADM [Adj]
    commonSOAC Pat Type
    pat StmAux ()
    aux SOAC SOACS
    soac ADM ()
    m = do
      Stm (Rep ADM) -> ADM ()
    forall (m :: * -> *). MonadBuilder m => Stm (Rep m) -> m ()
    addStm (Stm (Rep ADM) -> ADM ()) -> Stm (Rep ADM) -> ADM ()
    forall a b. (a -> b) -> a -> b
    $ Pat (LetDec (Rep ADM))
    -> StmAux (ExpDec (Rep ADM)) -> Exp (Rep ADM) -> Stm (Rep ADM)
    forall rep.
    Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
    Let Pat Type
    Pat (LetDec (Rep ADM))
    pat StmAux ()
    StmAux (ExpDec (Rep ADM))
    aux (Exp (Rep ADM) -> Stm (Rep ADM)) -> Exp (Rep ADM) -> Stm (Rep ADM)
    forall a b. (a -> b) -> a -> b
    $ Op SOACS -> Exp SOACS
    forall rep. Op rep -> Exp rep
    Op Op SOACS
    SOAC SOACS
    soac
      ADM ()
    m
      ADM [Adj] -> ADM [Adj]
    forall a. ADM a -> ADM a
    returnSweepCode (ADM [Adj] -> ADM [Adj]) -> ADM [Adj] -> ADM [Adj]
    forall a b. (a -> b) -> a -> b
    $ (VName -> ADM Adj) -> [VName] -> ADM [Adj]
    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 VName -> ADM Adj
    lookupAdj ([VName] -> ADM [Adj]) -> [VName] -> ADM [Adj]
    forall a b. (a -> b) -> a -> b
    $ Pat Type -> [VName]
    forall dec. Pat dec -> [VName]
    patNames Pat Type
    pat

    -- Reverse-mode differentiation of SOACs
    vjpSOAC :: VjpOps -> Pat Type -> StmAux () -> SOAC SOACS -> ADM () -> ADM ()
    -- Differentiating Reduces
    vjpSOAC :: VjpOps -> Pat Type -> StmAux () -> SOAC SOACS -> ADM () -> ADM ()
    vjpSOAC VjpOps
    ops Pat Type
    pat StmAux ()
    aux soac :: SOAC SOACS
    soac@(Screma SubExp
    w [VName]
    as ScremaForm SOACS
    form) ADM ()
    m
      | Just [Reduce Commutativity
    iscomm Lambda SOACS
    lam [Var VName
    ne]] <- ScremaForm SOACS -> Maybe [Reduce SOACS]
    forall rep. ScremaForm rep -> Maybe [Reduce rep]
    isReduceSOAC ScremaForm SOACS
    form,
        [VName
    a] <- [VName]
    as,
        Just Lambda SOACS
    op <- Lambda SOACS -> Maybe (Lambda SOACS)
    mapOp Lambda SOACS
    lam =
          VjpOps
    -> Pat Type
    -> StmAux ()
    -> SubExp
    -> Commutativity
    -> Lambda SOACS
    -> VName
    -> VName
    -> ADM ()
    -> ADM ()
    diffVecReduce VjpOps
    ops Pat Type
    pat StmAux ()
    aux SubExp
    w Commutativity
    iscomm Lambda SOACS
    op VName
    ne VName
    a ADM ()
    m
      | Just [Reduce SOACS]
    reds <- ScremaForm SOACS -> Maybe [Reduce SOACS]
    forall rep. ScremaForm rep -> Maybe [Reduce rep]
    isReduceSOAC ScremaForm SOACS
    form,
        [Reduce SOACS] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [Reduce SOACS]
    reds Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    > Int
    1 =
          VjpOps
    -> ([Reduce SOACS] -> ADM (ScremaForm SOACS),
        Reduce SOACS -> [SubExp])
    -> (Pat Type, StmAux (), [Reduce SOACS], SubExp, [VName])
    -> ADM ()
    -> ADM ()
    forall a.
    VjpOps
    -> ([a] -> ADM (ScremaForm SOACS), a -> [SubExp])
    -> (Pat Type, StmAux (), [a], SubExp, [VName])
    -> ADM ()
    -> ADM ()
    splitScanRed VjpOps
    ops ([Reduce SOACS] -> ADM (ScremaForm SOACS)
    forall rep (m :: * -> *).
    (Buildable rep, MonadFreshNames m) =>
    [Reduce rep] -> m (ScremaForm rep)
    reduceSOAC, Reduce SOACS -> [SubExp]
    forall rep. Reduce rep -> [SubExp]
    redNeutral) (Pat Type
    pat, StmAux ()
    aux, [Reduce SOACS]
    reds, SubExp
    w, [VName]
    as) ADM ()
    m
      | Just [Reduce SOACS
    red] <- ScremaForm SOACS -> Maybe [Reduce SOACS]
    forall rep. ScremaForm rep -> Maybe [Reduce rep]
    isReduceSOAC ScremaForm SOACS
    form,
        [VName
    x] <- Pat Type -> [VName]
    forall dec. Pat dec -> [VName]
    patNames Pat Type
    pat,
        [SubExp
    ne] <- Reduce SOACS -> [SubExp]
    forall rep. Reduce rep -> [SubExp]
    redNeutral Reduce SOACS
    red,
        [VName
    a] <- [VName]
    as,
        Just [(BinOp
    op, PrimType
    _, VName
    _, VName
    _)] <- Lambda SOACS -> Maybe [(BinOp, PrimType, VName, VName)]
    forall rep.
    ASTRep rep =>
    Lambda rep -> Maybe [(BinOp, PrimType, VName, VName)]
    lamIsBinOp (Lambda SOACS -> Maybe [(BinOp, PrimType, VName, VName)])
    -> Lambda SOACS -> Maybe [(BinOp, PrimType, VName, VName)]
    forall a b. (a -> b) -> a -> b
    $ Reduce SOACS -> Lambda SOACS
    forall rep. Reduce rep -> Lambda rep
    redLambda Reduce SOACS
    red,
        BinOp -> Bool
    isMinMaxOp BinOp
    op =
          VjpOps
    -> VName
    -> StmAux ()
    -> SubExp
    -> BinOp
    -> SubExp
    -> VName
    -> ADM ()
    -> ADM ()
    diffMinMaxReduce VjpOps
    ops VName
    x StmAux ()
    aux SubExp
    w BinOp
    op SubExp
    ne VName
    a ADM ()
    m
      | Just [Reduce SOACS
    red] <- ScremaForm SOACS -> Maybe [Reduce SOACS]
    forall rep. ScremaForm rep -> Maybe [Reduce rep]
    isReduceSOAC ScremaForm SOACS
    form,
        [VName
    x] <- Pat Type -> [VName]
    forall dec. Pat dec -> [VName]
    patNames Pat Type
    pat,
        [SubExp
    ne] <- Reduce SOACS -> [SubExp]
    forall rep. Reduce rep -> [SubExp]
    redNeutral Reduce SOACS
    red,
        [VName
    a] <- [VName]
    as,
        Just [(BinOp
    op, PrimType
    _, VName
    _, VName
    _)] <- Lambda SOACS -> Maybe [(BinOp, PrimType, VName, VName)]
    forall rep.
    ASTRep rep =>
    Lambda rep -> Maybe [(BinOp, PrimType, VName, VName)]
    lamIsBinOp (Lambda SOACS -> Maybe [(BinOp, PrimType, VName, VName)])
    -> Lambda SOACS -> Maybe [(BinOp, PrimType, VName, VName)]
    forall a b. (a -> b) -> a -> b
    $ Reduce SOACS -> Lambda SOACS
    forall rep. Reduce rep -> Lambda rep
    redLambda Reduce SOACS
    red,
        BinOp -> Bool
    isMulOp BinOp
    op =
          VjpOps
    -> VName
    -> StmAux ()
    -> SubExp
    -> BinOp
    -> SubExp
    -> VName
    -> ADM ()
    -> ADM ()
    diffMulReduce VjpOps
    ops VName
    x StmAux ()
    aux SubExp
    w BinOp
    op SubExp
    ne VName
    a ADM ()
    m
      | Just Reduce SOACS
    red <- [Reduce SOACS] -> Reduce SOACS
    forall rep. Buildable rep => [Reduce rep] -> Reduce rep
    singleReduce ([Reduce SOACS] -> Reduce SOACS)
    -> Maybe [Reduce SOACS] -> Maybe (Reduce SOACS)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> ScremaForm SOACS -> Maybe [Reduce SOACS]
    forall rep. ScremaForm rep -> Maybe [Reduce rep]
    isReduceSOAC ScremaForm SOACS
    form = do
          pat_adj <- (Adj -> ADM VName) -> [Adj] -> ADM [VName]
    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 Adj -> ADM VName
    adjVal ([Adj] -> ADM [VName]) -> ADM [Adj] -> ADM [VName]
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< Pat Type -> StmAux () -> SOAC SOACS -> ADM () -> ADM [Adj]
    commonSOAC Pat Type
    pat StmAux ()
    aux SOAC SOACS
    soac ADM ()
    m
          diffReduce ops pat_adj w as red

    -- Differentiating Scans
    vjpSOAC VjpOps
    ops Pat Type
    pat StmAux ()
    aux soac :: SOAC SOACS
    soac@(Screma SubExp
    w [VName]
    as ScremaForm SOACS
    form) ADM ()
    m
      | Just [Scan Lambda SOACS
    lam [SubExp
    ne]] <- ScremaForm SOACS -> Maybe [Scan SOACS]
    forall rep. ScremaForm rep -> Maybe [Scan rep]
    isScanSOAC ScremaForm SOACS
    form,
        [VName
    x] <- Pat Type -> [VName]
    forall dec. Pat dec -> [VName]
    patNames Pat Type
    pat,
        [VName
    a] <- [VName]
    as,
        Just [(BinOp
    op, PrimType
    _, VName
    _, VName
    _)] <- Lambda SOACS -> Maybe [(BinOp, PrimType, VName, VName)]
    forall rep.
    ASTRep rep =>
    Lambda rep -> Maybe [(BinOp, PrimType, VName, VName)]
    lamIsBinOp Lambda SOACS
    lam,
        BinOp -> Bool
    isAddOp BinOp
    op = do
          ADM [Adj] -> ADM ()
    forall (f :: * -> *) a. Functor f => f a -> f ()
    void (ADM [Adj] -> ADM ()) -> ADM [Adj] -> ADM ()
    forall a b. (a -> b) -> a -> b
    $ Pat Type -> StmAux () -> SOAC SOACS -> ADM () -> ADM [Adj]
    commonSOAC Pat Type
    pat StmAux ()
    aux SOAC SOACS
    soac ADM ()
    m
          VjpOps
    -> VName -> SubExp -> Lambda SOACS -> SubExp -> VName -> ADM ()
    diffScanAdd VjpOps
    ops VName
    x SubExp
    w Lambda SOACS
    lam SubExp
    ne VName
    a
      | Just [Scan Lambda SOACS
    lam [SubExp]
    ne] <- ScremaForm SOACS -> Maybe [Scan SOACS]
    forall rep. ScremaForm rep -> Maybe [Scan rep]
    isScanSOAC ScremaForm SOACS
    form,
        Just Lambda SOACS
    op <- Lambda SOACS -> Maybe (Lambda SOACS)
    mapOp Lambda SOACS
    lam = do
          VjpOps
    -> [VName]
    -> StmAux ()
    -> SubExp
    -> Lambda SOACS
    -> [SubExp]
    -> [VName]
    -> ADM ()
    -> ADM ()
    diffScanVec VjpOps
    ops (Pat Type -> [VName]
    forall dec. Pat dec -> [VName]
    patNames Pat Type
    pat) StmAux ()
    aux SubExp
    w Lambda SOACS
    op [SubExp]
    ne [VName]
    as ADM ()
    m
      | Just [Scan SOACS]
    scans <- ScremaForm SOACS -> Maybe [Scan SOACS]
    forall rep. ScremaForm rep -> Maybe [Scan rep]
    isScanSOAC ScremaForm SOACS
    form,
        [Scan SOACS] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [Scan SOACS]
    scans Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    > Int
    1 =
          VjpOps
    -> ([Scan SOACS] -> ADM (ScremaForm SOACS), Scan SOACS -> [SubExp])
    -> (Pat Type, StmAux (), [Scan SOACS], SubExp, [VName])
    -> ADM ()
    -> ADM ()
    forall a.
    VjpOps
    -> ([a] -> ADM (ScremaForm SOACS), a -> [SubExp])
    -> (Pat Type, StmAux (), [a], SubExp, [VName])
    -> ADM ()
    -> ADM ()
    splitScanRed VjpOps
    ops ([Scan SOACS] -> ADM (ScremaForm SOACS)
    forall rep (m :: * -> *).
    (Buildable rep, MonadFreshNames m) =>
    [Scan rep] -> m (ScremaForm rep)
    scanSOAC, Scan SOACS -> [SubExp]
    forall rep. Scan rep -> [SubExp]
    scanNeutral) (Pat Type
    pat, StmAux ()
    aux, [Scan SOACS]
    scans, SubExp
    w, [VName]
    as) ADM ()
    m
      | Just Scan SOACS
    red <- [Scan SOACS] -> Scan SOACS
    forall rep. Buildable rep => [Scan rep] -> Scan rep
    singleScan ([Scan SOACS] -> Scan SOACS)
    -> Maybe [Scan SOACS] -> Maybe (Scan SOACS)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> ScremaForm SOACS -> Maybe [Scan SOACS]
    forall rep. ScremaForm rep -> Maybe [Scan rep]
    isScanSOAC ScremaForm SOACS
    form = do
          ADM [Adj] -> ADM ()
    forall (f :: * -> *) a. Functor f => f a -> f ()
    void (ADM [Adj] -> ADM ()) -> ADM [Adj] -> ADM ()
    forall a b. (a -> b) -> a -> b
    $ Pat Type -> StmAux () -> SOAC SOACS -> ADM () -> ADM [Adj]
    commonSOAC Pat Type
    pat StmAux ()
    aux SOAC SOACS
    soac ADM ()
    m
          VjpOps -> [VName] -> SubExp -> [VName] -> Scan SOACS -> ADM ()
    diffScan VjpOps
    ops (Pat Type -> [VName]
    forall dec. Pat dec -> [VName]
    patNames Pat Type
    pat) SubExp
    w [VName]
    as Scan SOACS
    red

    -- Differentiating Maps
    vjpSOAC VjpOps
    ops Pat Type
    pat StmAux ()
    aux soac :: SOAC SOACS
    soac@(Screma SubExp
    w [VName]
    as ScremaForm SOACS
    form) ADM ()
    m
      | Just Lambda SOACS
    lam <- ScremaForm SOACS -> Maybe (Lambda SOACS)
    forall rep. ScremaForm rep -> Maybe (Lambda rep)
    isMapSOAC ScremaForm SOACS
    form = do
          pat_adj <- Pat Type -> StmAux () -> SOAC SOACS -> ADM () -> ADM [Adj]
    commonSOAC Pat Type
    pat StmAux ()
    aux SOAC SOACS
    soac ADM ()
    m
          vjpMap ops pat_adj aux w lam as

    -- Differentiating Redomaps
    vjpSOAC VjpOps
    ops Pat Type
    pat StmAux ()
    _aux (Screma SubExp
    w [VName]
    as ScremaForm SOACS
    form) ADM ()
    m
      | Just ([Reduce SOACS]
    reds, Lambda SOACS
    map_lam) <-
          ScremaForm SOACS -> Maybe ([Reduce SOACS], Lambda SOACS)
    forall rep. ScremaForm rep -> Maybe ([Reduce rep], Lambda rep)
    isRedomapSOAC ScremaForm SOACS
    form = do
          (mapstm, redstm) <-
            Pat (LetDec SOACS)
    -> (SubExp, [Reduce SOACS], Lambda SOACS, [VName])
    -> ADM (Stm SOACS, Stm SOACS)
    forall (m :: * -> *) rep.
    (MonadFreshNames m, Buildable rep, ExpDec rep ~ (),
     Op rep ~ SOAC rep) =>
    Pat (LetDec rep)
    -> (SubExp, [Reduce rep], Lambda rep, [VName])
    -> m (Stm rep, Stm rep)
    redomapToMapAndReduce Pat Type
    Pat (LetDec SOACS)
    pat (SubExp
    w, [Reduce SOACS]
    reds, Lambda SOACS
    map_lam, [VName]
    as)
          vjpStm ops mapstm $ vjpStm ops redstm m

    -- Differentiating Scanomaps
    vjpSOAC VjpOps
    ops Pat Type
    pat StmAux ()
    _aux (Screma SubExp
    w [VName]
    as ScremaForm SOACS
    form) ADM ()
    m
      | Just ([Scan SOACS]
    scans, Lambda SOACS
    map_lam) <-
          ScremaForm SOACS -> Maybe ([Scan SOACS], Lambda SOACS)
    forall rep. ScremaForm rep -> Maybe ([Scan rep], Lambda rep)
    isScanomapSOAC ScremaForm SOACS
    form = do
          (mapstm, scanstm) <-
            Pat (LetDec SOACS)
    -> (SubExp, [Scan SOACS], Lambda SOACS, [VName])
    -> ADM (Stm SOACS, Stm SOACS)
    forall (m :: * -> *) rep.
    (MonadFreshNames m, Buildable rep, ExpDec rep ~ (),
     Op rep ~ SOAC rep) =>
    Pat (LetDec rep)
    -> (SubExp, [Scan rep], Lambda rep, [VName])
    -> m (Stm rep, Stm rep)
    scanomapToMapAndScan Pat Type
    Pat (LetDec SOACS)
    pat (SubExp
    w, [Scan SOACS]
    scans, Lambda SOACS
    map_lam, [VName]
    as)
          vjpStm ops mapstm $ vjpStm ops scanstm m

    -- Differentiating Scatter
    vjpSOAC VjpOps
    ops Pat Type
    pat StmAux ()
    aux (Scatter SubExp
    w [VName]
    ass ScatterSpec VName
    written_info Lambda SOACS
    lam) ADM ()
    m
      | Lambda SOACS -> Bool
    forall rep. Lambda rep -> Bool
    isIdentityLambda Lambda SOACS
    lam =
          VjpOps
    -> Pat Type
    -> StmAux ()
    -> (SubExp, [VName], Lambda SOACS, ScatterSpec VName)
    -> ADM ()
    -> ADM ()
    vjpScatter VjpOps
    ops Pat Type
    pat StmAux ()
    aux (SubExp
    w, [VName]
    ass, Lambda SOACS
    lam, ScatterSpec VName
    written_info) ADM ()
    m
      | Bool
    otherwise = do
          map_idents <- (Type -> ADM Ident) -> [Type] -> ADM [Ident]
    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 (\Type
    t -> [Char] -> Type -> ADM Ident
    forall (m :: * -> *).
    MonadFreshNames m =>
    [Char] -> Type -> m Ident
    newIdent [Char]
    "map_res" (Type -> SubExp -> Type
    forall d.
    ArrayShape (ShapeBase d) =>
    TypeBase (ShapeBase d) NoUniqueness
    -> d -> TypeBase (ShapeBase d) NoUniqueness
    arrayOfRow Type
    t SubExp
    w)) ([Type] -> ADM [Ident]) -> [Type] -> ADM [Ident]
    forall a b. (a -> b) -> a -> b
    $ Lambda SOACS -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda SOACS
    lam
          let map_stm = [Ident] -> Exp SOACS -> Stm SOACS
    forall rep. Buildable rep => [Ident] -> Exp rep -> Stm rep
    mkLet [Ident]
    map_idents (Exp SOACS -> Stm SOACS) -> Exp SOACS -> Stm SOACS
    forall a b. (a -> b) -> a -> b
    $ Op SOACS -> Exp SOACS
    forall rep. Op rep -> Exp rep
    Op (Op SOACS -> Exp SOACS) -> Op SOACS -> Exp SOACS
    forall a b. (a -> b) -> a -> b
    $ SubExp -> [VName] -> ScremaForm SOACS -> SOAC SOACS
    forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
    Screma SubExp
    w [VName]
    ass (ScremaForm SOACS -> SOAC SOACS) -> ScremaForm SOACS -> SOAC SOACS
    forall a b. (a -> b) -> a -> b
    $ Lambda SOACS -> ScremaForm SOACS
    forall rep. Lambda rep -> ScremaForm rep
    mapSOAC Lambda SOACS
    lam
          lam_id <- mkIdentityLambda $ lambdaReturnType lam
          let scatter_stm = Pat (LetDec SOACS)
    -> StmAux (ExpDec SOACS) -> Exp SOACS -> Stm SOACS
    forall rep.
    Pat (LetDec rep) -> StmAux (ExpDec rep) -> Exp rep -> Stm rep
    Let Pat Type
    Pat (LetDec SOACS)
    pat StmAux ()
    StmAux (ExpDec SOACS)
    aux (Exp SOACS -> Stm SOACS) -> Exp SOACS -> Stm SOACS
    forall a b. (a -> b) -> a -> b
    $ Op SOACS -> Exp SOACS
    forall rep. Op rep -> Exp rep
    Op (Op SOACS -> Exp SOACS) -> Op SOACS -> Exp SOACS
    forall a b. (a -> b) -> a -> b
    $ SubExp
    -> [VName] -> ScatterSpec VName -> Lambda SOACS -> SOAC SOACS
    forall rep.
    SubExp -> [VName] -> ScatterSpec VName -> Lambda rep -> SOAC rep
    Scatter SubExp
    w ((Ident -> VName) -> [Ident] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map Ident -> VName
    identName [Ident]
    map_idents) ScatterSpec VName
    written_info Lambda SOACS
    lam_id
          vjpStm ops map_stm $ vjpStm ops scatter_stm m

    -- Differentiating Histograms
    vjpSOAC VjpOps
    ops Pat Type
    pat StmAux ()
    aux (Hist SubExp
    n [VName]
    as [HistOp SOACS]
    histops Lambda SOACS
    f) ADM ()
    m
      | Lambda SOACS -> Bool
    forall rep. Lambda rep -> Bool
    isIdentityLambda Lambda SOACS
    f,
        [HistOp SOACS] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [HistOp SOACS]
    histops Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    > Int
    1 = do
          let ([VName]
    is, [VName]
    vs) = Int -> [VName] -> ([VName], [VName])
    forall a. Int -> [a] -> ([a], [a])
    splitAt ([HistOp SOACS] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [HistOp SOACS]
    histops) [VName]
    as
          VjpOps
    -> Pat Type
    -> StmAux ()
    -> [HistOp SOACS]
    -> SubExp
    -> [VName]
    -> [VName]
    -> ADM ()
    -> ADM ()
    splitHist VjpOps
    ops Pat Type
    pat StmAux ()
    aux [HistOp SOACS]
    histops SubExp
    n [VName]
    is [VName]
    vs ADM ()
    m
    vjpSOAC VjpOps
    ops Pat Type
    pat StmAux ()
    aux (Hist SubExp
    n [VName
    is, VName
    vs] [HistOp SOACS
    histop] Lambda SOACS
    f) ADM ()
    m
      | Lambda SOACS -> Bool
    forall rep. Lambda rep -> Bool
    isIdentityLambda Lambda SOACS
    f,
        [VName
    x] <- Pat Type -> [VName]
    forall dec. Pat dec -> [VName]
    patNames Pat Type
    pat,
        HistOp (Shape [SubExp
    w]) SubExp
    rf [VName
    dst] [Var VName
    ne] Lambda SOACS
    lam <- HistOp SOACS
    histop,
        -- Note that the operator is vectorised, so `ne` cannot be a 'PrimValue'.
        Just Lambda SOACS
    op <- Lambda SOACS -> Maybe (Lambda SOACS)
    mapOp Lambda SOACS
    lam =
          VjpOps
    -> VName
    -> StmAux ()
    -> SubExp
    -> Lambda SOACS
    -> VName
    -> VName
    -> VName
    -> SubExp
    -> SubExp
    -> VName
    -> ADM ()
    -> ADM ()
    diffVecHist VjpOps
    ops VName
    x StmAux ()
    aux SubExp
    n Lambda SOACS
    op VName
    ne VName
    is VName
    vs SubExp
    w SubExp
    rf VName
    dst ADM ()
    m
      | Lambda SOACS -> Bool
    forall rep. Lambda rep -> Bool
    isIdentityLambda Lambda SOACS
    f,
        [VName
    x] <- Pat Type -> [VName]
    forall dec. Pat dec -> [VName]
    patNames Pat Type
    pat,
        HistOp (Shape [SubExp
    w]) SubExp
    rf [VName
    dst] [SubExp
    ne] Lambda SOACS
    lam <- HistOp SOACS
    histop,
        Lambda SOACS
    lam' <- Lambda SOACS -> Lambda SOACS
    nestedMapOp Lambda SOACS
    lam,
        Just [(BinOp
    op, PrimType
    _, VName
    _, VName
    _)] <- Lambda SOACS -> Maybe [(BinOp, PrimType, VName, VName)]
    forall rep.
    ASTRep rep =>
    Lambda rep -> Maybe [(BinOp, PrimType, VName, VName)]
    lamIsBinOp Lambda SOACS
    lam',
        BinOp -> Bool
    isMinMaxOp BinOp
    op =
          VjpOps
    -> VName
    -> StmAux ()
    -> SubExp
    -> BinOp
    -> SubExp
    -> VName
    -> VName
    -> SubExp
    -> SubExp
    -> VName
    -> ADM ()
    -> ADM ()
    diffMinMaxHist VjpOps
    ops VName
    x StmAux ()
    aux SubExp
    n BinOp
    op SubExp
    ne VName
    is VName
    vs SubExp
    w SubExp
    rf VName
    dst ADM ()
    m
      | Lambda SOACS -> Bool
    forall rep. Lambda rep -> Bool
    isIdentityLambda Lambda SOACS
    f,
        [VName
    x] <- Pat Type -> [VName]
    forall dec. Pat dec -> [VName]
    patNames Pat Type
    pat,
        HistOp (Shape [SubExp
    w]) SubExp
    rf [VName
    dst] [SubExp
    ne] Lambda SOACS
    lam <- HistOp SOACS
    histop,
        Lambda SOACS
    lam' <- Lambda SOACS -> Lambda SOACS
    nestedMapOp Lambda SOACS
    lam,
        Just [(BinOp
    op, PrimType
    _, VName
    _, VName
    _)] <- Lambda SOACS -> Maybe [(BinOp, PrimType, VName, VName)]
    forall rep.
    ASTRep rep =>
    Lambda rep -> Maybe [(BinOp, PrimType, VName, VName)]
    lamIsBinOp Lambda SOACS
    lam',
        BinOp -> Bool
    isMulOp BinOp
    op =
          VjpOps
    -> VName
    -> StmAux ()
    -> SubExp
    -> BinOp
    -> SubExp
    -> VName
    -> VName
    -> SubExp
    -> SubExp
    -> VName
    -> ADM ()
    -> ADM ()
    diffMulHist VjpOps
    ops VName
    x StmAux ()
    aux SubExp
    n BinOp
    op SubExp
    ne VName
    is VName
    vs SubExp
    w SubExp
    rf VName
    dst ADM ()
    m
      | Lambda SOACS -> Bool
    forall rep. Lambda rep -> Bool
    isIdentityLambda Lambda SOACS
    f,
        [VName
    x] <- Pat Type -> [VName]
    forall dec. Pat dec -> [VName]
    patNames Pat Type
    pat,
        HistOp (Shape [SubExp
    w]) SubExp
    rf [VName
    dst] [SubExp
    ne] Lambda SOACS
    lam <- HistOp SOACS
    histop,
        Lambda SOACS
    lam' <- Lambda SOACS -> Lambda SOACS
    nestedMapOp Lambda SOACS
    lam,
        Just [(BinOp
    op, PrimType
    _, VName
    _, VName
    _)] <- Lambda SOACS -> Maybe [(BinOp, PrimType, VName, VName)]
    forall rep.
    ASTRep rep =>
    Lambda rep -> Maybe [(BinOp, PrimType, VName, VName)]
    lamIsBinOp Lambda SOACS
    lam',
        BinOp -> Bool
    isAddOp BinOp
    op =
          VjpOps
    -> VName
    -> StmAux ()
    -> SubExp
    -> Lambda SOACS
    -> SubExp
    -> VName
    -> VName
    -> SubExp
    -> SubExp
    -> VName
    -> ADM ()
    -> ADM ()
    diffAddHist VjpOps
    ops VName
    x StmAux ()
    aux SubExp
    n Lambda SOACS
    lam SubExp
    ne VName
    is VName
    vs SubExp
    w SubExp
    rf VName
    dst ADM ()
    m
    vjpSOAC VjpOps
    ops Pat Type
    pat StmAux ()
    aux (Hist SubExp
    n [VName]
    as [HistOp SOACS
    histop] Lambda SOACS
    f) ADM ()
    m
      | Lambda SOACS -> Bool
    forall rep. Lambda rep -> Bool
    isIdentityLambda Lambda SOACS
    f,
        HistOp (Shape [SubExp]
    w) SubExp
    rf [VName]
    dst [SubExp]
    ne Lambda SOACS
    lam <- HistOp SOACS
    histop = do
          VjpOps
    -> [VName]
    -> StmAux ()
    -> SubExp
    -> Lambda SOACS
    -> [SubExp]
    -> [VName]
    -> [SubExp]
    -> SubExp
    -> [VName]
    -> ADM ()
    -> ADM ()
    diffHist VjpOps
    ops (Pat Type -> [VName]
    forall dec. Pat dec -> [VName]
    patNames Pat Type
    pat) StmAux ()
    aux SubExp
    n Lambda SOACS
    lam [SubExp]
    ne [VName]
    as [SubExp]
    w SubExp
    rf [VName]
    dst ADM ()
    m
    vjpSOAC VjpOps
    ops Pat Type
    pat StmAux ()
    _aux (Hist SubExp
    n [VName]
    as [HistOp SOACS]
    histops Lambda SOACS
    f) ADM ()
    m
      | Bool -> Bool
    not (Lambda SOACS -> Bool
    forall rep. Lambda rep -> Bool
    isIdentityLambda Lambda SOACS
    f) = do
          (mapstm, redstm) <-
            Pat Type
    -> (SubExp, [HistOp SOACS], Lambda SOACS, [VName])
    -> ADM (Stm SOACS, Stm SOACS)
    histomapToMapAndHist Pat Type
    pat (SubExp
    n, [HistOp SOACS]
    histops, Lambda SOACS
    f, [VName]
    as)
          vjpStm ops mapstm $ vjpStm ops redstm m
    vjpSOAC VjpOps
    ops Pat Type
    pat StmAux ()
    aux (Stream SubExp
    w [VName]
    as [SubExp]
    accs Lambda SOACS
    lam) ADM ()
    m = do
      stms <- ADM () -> ADM (Stms (Rep ADM))
    forall (m :: * -> *) a. MonadBuilder m => m a -> m (Stms (Rep m))
    collectStms_ (ADM () -> ADM (Stms (Rep ADM))) -> ADM () -> ADM (Stms (Rep ADM))
    forall a b. (a -> b) -> a -> b
    $ StmAux () -> ADM () -> ADM ()
    forall (m :: * -> *) anyrep a.
    MonadBuilder m =>
    StmAux anyrep -> m a -> m a
    auxing StmAux ()
    aux (ADM () -> ADM ()) -> ADM () -> ADM ()
    forall a b. (a -> b) -> a -> b
    $ Pat (LetDec (Rep ADM))
    -> SubExp -> [SubExp] -> Lambda (Rep ADM) -> [VName] -> ADM ()
    forall (m :: * -> *).
    (MonadBuilder m, Buildable (Rep m)) =>
    Pat (LetDec (Rep m))
    -> SubExp -> [SubExp] -> Lambda (Rep m) -> [VName] -> m ()
    sequentialStreamWholeArray Pat Type
    Pat (LetDec (Rep ADM))
    pat SubExp
    w [SubExp]
    accs Lambda (Rep ADM)
    Lambda SOACS
    lam [VName]
    as
      foldr (vjpStm ops) m stms
    vjpSOAC VjpOps
    _ Pat Type
    _ StmAux ()
    _ SOAC SOACS
    soac ADM ()
    _ =
      [Char] -> ADM ()
    forall a. HasCallStack => [Char] -> a
    error ([Char] -> ADM ()) -> [Char] -> ADM ()
    forall a b. (a -> b) -> a -> b
    $ [Char]
    "vjpSOAC unhandled:\n" [Char] -> [Char] -> [Char]
    forall a. [a] -> [a] -> [a]
    ++ SOAC SOACS -> [Char]
    forall a. Pretty a => a -> [Char]
    prettyString SOAC SOACS
    soac

    ---------------
    --- Helpers ---
    ---------------

    isMinMaxOp :: BinOp -> Bool
    isMinMaxOp :: BinOp -> Bool
    isMinMaxOp (SMin IntType
    _) = Bool
    True
    isMinMaxOp (UMin IntType
    _) = Bool
    True
    isMinMaxOp (FMin FloatType
    _) = Bool
    True
    isMinMaxOp (SMax IntType
    _) = Bool
    True
    isMinMaxOp (UMax IntType
    _) = Bool
    True
    isMinMaxOp (FMax FloatType
    _) = Bool
    True
    isMinMaxOp BinOp
    _ = Bool
    False

    isMulOp :: BinOp -> Bool
    isMulOp :: BinOp -> Bool
    isMulOp (Mul IntType
    _ Overflow
    _) = Bool
    True
    isMulOp (FMul FloatType
    _) = Bool
    True
    isMulOp BinOp
    _ = Bool
    False

    isAddOp :: BinOp -> Bool
    isAddOp :: BinOp -> Bool
    isAddOp (Add IntType
    _ Overflow
    _) = Bool
    True
    isAddOp (FAdd FloatType
    _) = Bool
    True
    isAddOp BinOp
    _ = Bool
    False

    -- Identifies vectorized operators (lambdas):
    --   if the lambda argument is a map, then returns
    --   just the map's lambda; otherwise nothing.
    mapOp :: Lambda SOACS -> Maybe (Lambda SOACS)
    mapOp :: Lambda SOACS -> Maybe (Lambda SOACS)
    mapOp (Lambda [LParam SOACS
    pa1, LParam SOACS
    pa2] [Type]
    _ Body SOACS
    lam_body)
      | [SubExpRes Certs
    cs SubExp
    r] <- Body SOACS -> [SubExpRes]
    forall rep. Body rep -> [SubExpRes]
    bodyResult Body SOACS
    lam_body,
        Certs
    cs Certs -> Certs -> Bool
    forall a. Eq a => a -> a -> Bool
    == Certs
    forall a. Monoid a => a
    mempty,
        [Stm SOACS
    map_stm] <- Seq (Stm SOACS) -> [Stm SOACS]
    forall rep. Stms rep -> [Stm rep]
    stmsToList (Body SOACS -> Seq (Stm SOACS)
    forall rep. Body rep -> Stms rep
    bodyStms Body SOACS
    lam_body),
        (Let (Pat [PatElem (LetDec SOACS)
    pe]) StmAux (ExpDec SOACS)
    _ (Op Op SOACS
    scrm)) <- Stm SOACS
    map_stm,
        (Screma SubExp
    _ [VName
    a1, VName
    a2] (ScremaForm Lambda SOACS
    map_lam [] [])) <- Op SOACS
    scrm,
        (VName
    a1 VName -> VName -> Bool
    forall a. Eq a => a -> a -> Bool
    == Param Type -> VName
    forall dec. Param dec -> VName
    paramName Param Type
    LParam SOACS
    pa1 Bool -> Bool -> Bool
    && VName
    a2 VName -> VName -> Bool
    forall a. Eq a => a -> a -> Bool
    == Param Type -> VName
    forall dec. Param dec -> VName
    paramName Param Type
    LParam SOACS
    pa2) Bool -> Bool -> Bool
    || (VName
    a1 VName -> VName -> Bool
    forall a. Eq a => a -> a -> Bool
    == Param Type -> VName
    forall dec. Param dec -> VName
    paramName Param Type
    LParam SOACS
    pa2 Bool -> Bool -> Bool
    && VName
    a2 VName -> VName -> Bool
    forall a. Eq a => a -> a -> Bool
    == Param Type -> VName
    forall dec. Param dec -> VName
    paramName Param Type
    LParam SOACS
    pa1),
        SubExp
    r SubExp -> SubExp -> Bool
    forall a. Eq a => a -> a -> Bool
    == VName -> SubExp
    Var (PatElem Type -> VName
    forall dec. PatElem dec -> VName
    patElemName PatElem Type
    PatElem (LetDec SOACS)
    pe) =
          Lambda SOACS -> Maybe (Lambda SOACS)
    forall a. a -> Maybe a
    Just Lambda SOACS
    map_lam
    mapOp Lambda SOACS
    _ = Maybe (Lambda SOACS)
    forall a. Maybe a
    Nothing

    -- getting the innermost lambda of a perfect-map nest
    --   (i.e., the first lambda that does not consists of exactly a map)
    nestedMapOp :: Lambda SOACS -> Lambda SOACS
    nestedMapOp :: Lambda SOACS -> Lambda SOACS
    nestedMapOp Lambda SOACS
    lam =
      Lambda SOACS
    -> (Lambda SOACS -> Lambda SOACS)
    -> Maybe (Lambda SOACS)
    -> Lambda SOACS
    forall b a. b -> (a -> b) -> Maybe a -> b
    maybe Lambda SOACS
    lam Lambda SOACS -> Lambda SOACS
    nestedMapOp (Lambda SOACS -> Maybe (Lambda SOACS)
    mapOp Lambda SOACS
    lam)