    {-# LANGUAGE Strict #-}

    -- | This module consists of rules for fusion
    --     that involves WithAcc constructs.
    --   Currently, we support two non-trivial
    --   transformations:
    --     I. map-flatten-scatter: a map nest produces
    --        multi-dimensional index and values arrays
    --        that are then flattened and used in a
    --        scatter consumer. Such pattern can be fused
    --        by re-writing the scatter by means of a WithAcc
    --        containing a map-nest, thus eliminating the flatten
    --        operations. The obtained WithAcc can then be fused
    --        with the producer map nest, e.g., benefiting intra-group
    --        kernels. The eloquent target for this rule is
    --        an efficient implementation of radix-sort.
    --
    --    II. WithAcc-WithAcc fusion: two withaccs can be
    --        fused as long as the common accumulators use
    --        the same operator, and as long as the non-accumulator
    --        input of an WithAcc is not used as an accumulator in
    --        the other. This fusion opens the door for fusing
    --        the SOACs appearing inside the WithAccs. This is
    --        also intended to demonstrate that it is not so
    --        important where exactly the WithAccs were originally
    --        introduced in the code, it is more important that
    --        they can be transformed by various optimizations passes.
    module Futhark.Optimise.Fusion.RulesWithAccs
      ( ruleMFScat,
        tryFuseWithAccs,
      )
    where

    import Control.Monad
    import Data.Graph.Inductive.Graph qualified as G
    import Data.Map.Strict qualified as M
    import Data.Maybe
    import Futhark.Analysis.HORep.SOAC qualified as H
    import Futhark.Construct
    import Futhark.IR.SOACS hiding (SOAC (..))
    import Futhark.IR.SOACS qualified as F
    import Futhark.Optimise.Fusion.GraphRep
    import Futhark.Tools
    import Futhark.Transform.Rename
    import Futhark.Transform.Substitute

    se0 :: SubExp
    se0 :: SubExp
    se0 = IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    0

    se1 :: SubExp
    se1 :: SubExp
    se1 = IntType -> Integer -> SubExp
    intConst IntType
    Int64 Integer
    1

    -------------------------------------
    --- I. Map-Flatten-Scatter Fusion ---
    -------------------------------------

    -- helper data structures
    type IotaInp = ((VName, LParam SOACS), (SubExp, SubExp, SubExp, IntType))
    -- ^           ((array-name, lambda param), (len, start, stride, Int64))

    type RshpInp = ((VName, LParam SOACS), (Shape, Shape, Type))
    -- ^           ((array-name, lambda param), (flat-shape, unflat-shape, elem-type))

    -- | Implements a specialized rule for fusing a pattern
    --   formed by a map o flatten o scatter, i.e.,
    --      let (inds,   vals) = map-nest f inps
    --          (finds, fvals) = (flatten inds, flatten vals)
    --          let res = scatter res0 finds fvals
    --   where inds & vals have higher rank than finds & fvals.
    ruleMFScat ::
      (HasScope SOACS m, MonadFreshNames m) =>
      DepNode ->
      DepGraph ->
      m (Maybe DepGraph)
    ruleMFScat :: forall (m :: * -> *).
    (HasScope SOACS m, MonadFreshNames m) =>
    DepNode -> DepGraph -> m (Maybe DepGraph)
    ruleMFScat DepNode
    node_to_fuse dg :: DepGraph
    dg@DepGraph {dgGraph :: DepGraph -> Gr NodeT EdgeT
    dgGraph = Gr NodeT EdgeT
    g}
      | NodeT
    soac_nodeT <- DepNode -> NodeT
    forall a b. (a, b) -> b
    snd DepNode
    node_to_fuse,
        Int
    scat_node_id <- DepNode -> Int
    nodeFromLNode DepNode
    node_to_fuse,
        SoacNode ArrayTransforms
    node_out_trsfs Pat Type
    scat_pat SOAC SOACS
    scat_soac StmAux (ExpDec SOACS)
    scat_aux <- NodeT
    soac_nodeT,
        ArrayTransforms -> Bool
    H.nullTransforms ArrayTransforms
    node_out_trsfs,
        -- \^ for simplicity we do not allow transforms on scatter's result.
        H.Scatter SubExp
    _len [Input]
    scat_inp ScatterSpec VName
    scat_out Lambda SOACS
    scat_lam <- SOAC SOACS
    scat_soac,
        -- \^ get the scatter
        [ArrayTransforms]
    scat_trsfs <- (Input -> ArrayTransforms) -> [Input] -> [ArrayTransforms]
    forall a b. (a -> b) -> [a] -> [b]
    map Input -> ArrayTransforms
    H.inputTransforms (SOAC SOACS -> [Input]
    forall rep. SOAC rep -> [Input]
    H.inputs SOAC SOACS
    scat_soac),
        -- \^ get the transforms on the input
        (ArrayTransforms -> Bool) -> [ArrayTransforms] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    any (ArrayTransforms -> ArrayTransforms -> Bool
    forall a. Eq a => a -> a -> Bool
    /= ArrayTransforms
    forall a. Monoid a => a
    mempty) [ArrayTransforms]
    scat_trsfs,
        Context NodeT EdgeT
    scat_ctx <- Gr NodeT EdgeT -> Int -> Context NodeT EdgeT
    forall (gr :: * -> * -> *) a b.
    Graph gr =>
    gr a b -> Int -> Context a b
    G.context Gr NodeT EdgeT
    g Int
    scat_node_id,
        (Adj EdgeT
    out_deps, Int
    _, NodeT
    _, Adj EdgeT
    inp_deps) <- Context NodeT EdgeT
    scat_ctx,
        Adj EdgeT
    cons_deps <- ((EdgeT, Int) -> Bool) -> Adj EdgeT -> Adj EdgeT
    forall a. (a -> Bool) -> [a] -> [a]
    filter (EdgeT -> Bool
    isCons (EdgeT -> Bool) -> ((EdgeT, Int) -> EdgeT) -> (EdgeT, Int) -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (EdgeT, Int) -> EdgeT
    forall a b. (a, b) -> a
    fst) Adj EdgeT
    inp_deps,
        Adj EdgeT
    drct_deps <- ((EdgeT, Int) -> Bool) -> Adj EdgeT -> Adj EdgeT
    forall a. (a -> Bool) -> [a] -> [a]
    filter (EdgeT -> Bool
    isDep (EdgeT -> Bool) -> ((EdgeT, Int) -> EdgeT) -> (EdgeT, Int) -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (EdgeT, Int) -> EdgeT
    forall a b. (a, b) -> a
    fst) Adj EdgeT
    inp_deps,
        [Context NodeT EdgeT]
    cons_ctxs <- ((EdgeT, Int) -> Context NodeT EdgeT)
    -> Adj EdgeT -> [Context NodeT EdgeT]
    forall a b. (a -> b) -> [a] -> [b]
    map (Gr NodeT EdgeT -> Int -> Context NodeT EdgeT
    forall (gr :: * -> * -> *) a b.
    Graph gr =>
    gr a b -> Int -> Context a b
    G.context Gr NodeT EdgeT
    g (Int -> Context NodeT EdgeT)
    -> ((EdgeT, Int) -> Int) -> (EdgeT, Int) -> Context NodeT EdgeT
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (EdgeT, Int) -> Int
    forall a b. (a, b) -> b
    snd) Adj EdgeT
    cons_deps,
        [Context NodeT EdgeT]
    drct_ctxs <- ((EdgeT, Int) -> Context NodeT EdgeT)
    -> Adj EdgeT -> [Context NodeT EdgeT]
    forall a b. (a -> b) -> [a] -> [b]
    map (Gr NodeT EdgeT -> Int -> Context NodeT EdgeT
    forall (gr :: * -> * -> *) a b.
    Graph gr =>
    gr a b -> Int -> Context a b
    G.context Gr NodeT EdgeT
    g (Int -> Context NodeT EdgeT)
    -> ((EdgeT, Int) -> Int) -> (EdgeT, Int) -> Context NodeT EdgeT
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (EdgeT, Int) -> Int
    forall a b. (a, b) -> b
    snd) Adj EdgeT
    drct_deps,
        [NodeT]
    _cons_nTs <- (Context NodeT EdgeT -> NodeT) -> [Context NodeT EdgeT] -> [NodeT]
    forall a b. (a -> b) -> [a] -> [b]
    map Context NodeT EdgeT -> NodeT
    forall {a} {b} {c} {d}. (a, b, c, d) -> c
    getNodeTfromCtx [Context NodeT EdgeT]
    cons_ctxs, -- not used!!
        [(Context NodeT EdgeT, (Input, NodeT))]
    drct_tups0 <- (Input -> Maybe (Context NodeT EdgeT, (Input, NodeT)))
    -> [Input] -> [(Context NodeT EdgeT, (Input, NodeT))]
    forall a b. (a -> Maybe b) -> [a] -> [b]
    mapMaybe ([(Context NodeT EdgeT, EdgeT)]
    -> Input -> Maybe (Context NodeT EdgeT, (Input, NodeT))
    pairUp ([Context NodeT EdgeT] -> [EdgeT] -> [(Context NodeT EdgeT, EdgeT)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Context NodeT EdgeT]
    drct_ctxs (((EdgeT, Int) -> EdgeT) -> Adj EdgeT -> [EdgeT]
    forall a b. (a -> b) -> [a] -> [b]
    map (EdgeT, Int) -> EdgeT
    forall a b. (a, b) -> a
    fst Adj EdgeT
    drct_deps))) [Input]
    scat_inp,
        [(Context NodeT EdgeT, (Input, NodeT))] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [(Context NodeT EdgeT, (Input, NodeT))]
    drct_tups0 Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == [Input] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [Input]
    scat_inp,
        -- \^ checks that all direct dependencies are also array
        --   inputs to scatter
        ([Context NodeT EdgeT]
    t1s, [(Input, NodeT)]
    t2s) <- [(Context NodeT EdgeT, (Input, NodeT))]
    -> ([Context NodeT EdgeT], [(Input, NodeT)])
    forall a b. [(a, b)] -> ([a], [b])
    unzip [(Context NodeT EdgeT, (Input, NodeT))]
    drct_tups0,
        [(Context NodeT EdgeT, ((Input, NodeT), Param Type))]
    drct_tups <- [Context NodeT EdgeT]
    -> [((Input, NodeT), Param Type)]
    -> [(Context NodeT EdgeT, ((Input, NodeT), Param Type))]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Context NodeT EdgeT]
    t1s ([((Input, NodeT), Param Type)]
     -> [(Context NodeT EdgeT, ((Input, NodeT), Param Type))])
    -> [((Input, NodeT), Param Type)]
    -> [(Context NodeT EdgeT, ((Input, NodeT), Param Type))]
    forall a b. (a -> b) -> a -> b
    $ [(Input, NodeT)] -> [Param Type] -> [((Input, NodeT), Param Type)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [(Input, NodeT)]
    t2s (Lambda SOACS -> [LParam SOACS]
    forall rep. Lambda rep -> [LParam rep]
    lambdaParams Lambda SOACS
    scat_lam),
        ([Context NodeT EdgeT]
    ctxs_iots, [((Input, NodeT), Param Type)]
    drct_iots) <- [(Context NodeT EdgeT, ((Input, NodeT), Param Type))]
    -> ([Context NodeT EdgeT], [((Input, NodeT), Param Type)])
    forall a b. [(a, b)] -> ([a], [b])
    unzip ([(Context NodeT EdgeT, ((Input, NodeT), Param Type))]
     -> ([Context NodeT EdgeT], [((Input, NodeT), Param Type)]))
    -> [(Context NodeT EdgeT, ((Input, NodeT), Param Type))]
    -> ([Context NodeT EdgeT], [((Input, NodeT), Param Type)])
    forall a b. (a -> b) -> a -> b
    $ ((Context NodeT EdgeT, ((Input, NodeT), Param Type)) -> Bool)
    -> [(Context NodeT EdgeT, ((Input, NodeT), Param Type))]
    -> [(Context NodeT EdgeT, ((Input, NodeT), Param Type))]
    forall a. (a -> Bool) -> [a] -> [a]
    filter (NodeT -> Bool
    isIota (NodeT -> Bool)
    -> ((Context NodeT EdgeT, ((Input, NodeT), Param Type)) -> NodeT)
    -> (Context NodeT EdgeT, ((Input, NodeT), Param Type))
    -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Input, NodeT) -> NodeT
    forall a b. (a, b) -> b
    snd ((Input, NodeT) -> NodeT)
    -> ((Context NodeT EdgeT, ((Input, NodeT), Param Type))
        -> (Input, NodeT))
    -> (Context NodeT EdgeT, ((Input, NodeT), Param Type))
    -> NodeT
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ((Input, NodeT), Param Type) -> (Input, NodeT)
    forall a b. (a, b) -> a
    fst (((Input, NodeT), Param Type) -> (Input, NodeT))
    -> ((Context NodeT EdgeT, ((Input, NodeT), Param Type))
        -> ((Input, NodeT), Param Type))
    -> (Context NodeT EdgeT, ((Input, NodeT), Param Type))
    -> (Input, NodeT)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Context NodeT EdgeT, ((Input, NodeT), Param Type))
    -> ((Input, NodeT), Param Type)
    forall a b. (a, b) -> b
    snd) [(Context NodeT EdgeT, ((Input, NodeT), Param Type))]
    drct_tups,
        ([Context NodeT EdgeT]
    ctxs_rshp, [((Input, NodeT), Param Type)]
    drct_rshp) <- [(Context NodeT EdgeT, ((Input, NodeT), Param Type))]
    -> ([Context NodeT EdgeT], [((Input, NodeT), Param Type)])
    forall a b. [(a, b)] -> ([a], [b])
    unzip ([(Context NodeT EdgeT, ((Input, NodeT), Param Type))]
     -> ([Context NodeT EdgeT], [((Input, NodeT), Param Type)]))
    -> [(Context NodeT EdgeT, ((Input, NodeT), Param Type))]
    -> ([Context NodeT EdgeT], [((Input, NodeT), Param Type)])
    forall a b. (a -> b) -> a -> b
    $ ((Context NodeT EdgeT, ((Input, NodeT), Param Type)) -> Bool)
    -> [(Context NodeT EdgeT, ((Input, NodeT), Param Type))]
    -> [(Context NodeT EdgeT, ((Input, NodeT), Param Type))]
    forall a. (a -> Bool) -> [a] -> [a]
    filter (Bool -> Bool
    not (Bool -> Bool)
    -> ((Context NodeT EdgeT, ((Input, NodeT), Param Type)) -> Bool)
    -> (Context NodeT EdgeT, ((Input, NodeT), Param Type))
    -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . NodeT -> Bool
    isIota (NodeT -> Bool)
    -> ((Context NodeT EdgeT, ((Input, NodeT), Param Type)) -> NodeT)
    -> (Context NodeT EdgeT, ((Input, NodeT), Param Type))
    -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Input, NodeT) -> NodeT
    forall a b. (a, b) -> b
    snd ((Input, NodeT) -> NodeT)
    -> ((Context NodeT EdgeT, ((Input, NodeT), Param Type))
        -> (Input, NodeT))
    -> (Context NodeT EdgeT, ((Input, NodeT), Param Type))
    -> NodeT
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ((Input, NodeT), Param Type) -> (Input, NodeT)
    forall a b. (a, b) -> a
    fst (((Input, NodeT), Param Type) -> (Input, NodeT))
    -> ((Context NodeT EdgeT, ((Input, NodeT), Param Type))
        -> ((Input, NodeT), Param Type))
    -> (Context NodeT EdgeT, ((Input, NodeT), Param Type))
    -> (Input, NodeT)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (Context NodeT EdgeT, ((Input, NodeT), Param Type))
    -> ((Input, NodeT), Param Type)
    forall a b. (a, b) -> b
    snd) [(Context NodeT EdgeT, ((Input, NodeT), Param Type))]
    drct_tups,
        [((Input, NodeT), Param Type)] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [((Input, NodeT), Param Type)]
    drct_iots Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + [((Input, NodeT), Param Type)] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [((Input, NodeT), Param Type)]
    drct_rshp Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == [Input] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [Input]
    scat_inp,
        -- \^ direct dependencies are either flatten reshapes or iotas.
        [((VName, Param Type), (SubExp, SubExp, SubExp, IntType))]
    rep_iotas <- (((Input, NodeT), Param Type)
     -> Maybe ((VName, Param Type), (SubExp, SubExp, SubExp, IntType)))
    -> [((Input, NodeT), Param Type)]
    -> [((VName, Param Type), (SubExp, SubExp, SubExp, IntType))]
    forall a b. (a -> Maybe b) -> [a] -> [b]
    mapMaybe ((Input, NodeT), Param Type)
    -> Maybe ((VName, Param Type), (SubExp, SubExp, SubExp, IntType))
    ((Input, NodeT), LParam SOACS) -> Maybe IotaInp
    getRepIota [((Input, NodeT), Param Type)]
    drct_iots,
        [((VName, Param Type), (SubExp, SubExp, SubExp, IntType))] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [((VName, Param Type), (SubExp, SubExp, SubExp, IntType))]
    rep_iotas Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == [((Input, NodeT), Param Type)] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [((Input, NodeT), Param Type)]
    drct_iots,
        [(((VName, Param Type),
       (ShapeBase SubExp, ShapeBase SubExp, Type)),
      Certs)]
    rep_rshps_certs <- (((Input, NodeT), Param Type)
     -> Maybe
          (((VName, Param Type), (ShapeBase SubExp, ShapeBase SubExp, Type)),
           Certs))
    -> [((Input, NodeT), Param Type)]
    -> [(((VName, Param Type),
          (ShapeBase SubExp, ShapeBase SubExp, Type)),
         Certs)]
    forall a b. (a -> Maybe b) -> [a] -> [b]
    mapMaybe ((Input, NodeT), Param Type)
    -> Maybe
         (((VName, Param Type), (ShapeBase SubExp, ShapeBase SubExp, Type)),
          Certs)
    ((Input, NodeT), LParam SOACS) -> Maybe (RshpInp, Certs)
    getRepRshpArr [((Input, NodeT), Param Type)]
    drct_rshp,
        ([((VName, Param Type), (ShapeBase SubExp, ShapeBase SubExp, Type))]
    rep_rshps, [Certs]
    certs_rshps) <- [(((VName, Param Type),
       (ShapeBase SubExp, ShapeBase SubExp, Type)),
      Certs)]
    -> ([((VName, Param Type),
          (ShapeBase SubExp, ShapeBase SubExp, Type))],
        [Certs])
    forall a b. [(a, b)] -> ([a], [b])
    unzip [(((VName, Param Type),
       (ShapeBase SubExp, ShapeBase SubExp, Type)),
      Certs)]
    rep_rshps_certs,
        -- \^ gather the representations for the iotas and reshapes, that use
        --   the helper types `IotaInp` and `RshpInp`
        Bool -> Bool
    not ([((VName, Param Type), (ShapeBase SubExp, ShapeBase SubExp, Type))]
    -> Bool
    forall a. [a] -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null [((VName, Param Type), (ShapeBase SubExp, ShapeBase SubExp, Type))]
    rep_rshps),
        -- \^ at least one flatten-reshaped array
        [((VName, Param Type), (ShapeBase SubExp, ShapeBase SubExp, Type))]
    -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [((VName, Param Type), (ShapeBase SubExp, ShapeBase SubExp, Type))]
    rep_rshps Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == [((Input, NodeT), Param Type)] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [((Input, NodeT), Param Type)]
    drct_rshp,
        ((VName, Param Type)
    _, (ShapeBase SubExp
    s1, ShapeBase SubExp
    s2, Type
    _)) : [((VName, Param Type), (ShapeBase SubExp, ShapeBase SubExp, Type))]
    _ <- [((VName, Param Type), (ShapeBase SubExp, ShapeBase SubExp, Type))]
    rep_rshps,
        (((VName, Param Type), (ShapeBase SubExp, ShapeBase SubExp, Type))
     -> Bool)
    -> [((VName, Param Type),
         (ShapeBase SubExp, ShapeBase SubExp, Type))]
    -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all (\((VName, Param Type)
    _, (ShapeBase SubExp
    s1', ShapeBase SubExp
    s2', Type
    _)) -> ShapeBase SubExp
    s1 ShapeBase SubExp -> ShapeBase SubExp -> Bool
    forall a. Eq a => a -> a -> Bool
    == ShapeBase SubExp
    s1' Bool -> Bool -> Bool
    && ShapeBase SubExp
    s2 ShapeBase SubExp -> ShapeBase SubExp -> Bool
    forall a. Eq a => a -> a -> Bool
    == ShapeBase SubExp
    s2') [((VName, Param Type), (ShapeBase SubExp, ShapeBase SubExp, Type))]
    rep_rshps,
        -- \^ Check that all unflatten shape dimensions are the same,
        --   so that we can construct a map nest;

        -- check profitability, which is conservatively defined as all
        -- the reshaped and consumer arrays are used solely by the
        -- scatter AND all reshape dependencies originate in the same
        -- map.
        DepGraph
    -> Int -> [Context NodeT EdgeT] -> [Context NodeT EdgeT] -> Bool
    checkSafeAndProfitable DepGraph
    dg Int
    scat_node_id [Context NodeT EdgeT]
    ctxs_rshp [Context NodeT EdgeT]
    cons_ctxs = do
          -- generate the withAcc statement
          let cons_patels_outs :: [(PatElem Type, (ShapeBase SubExp, Int, VName))]
    cons_patels_outs = [PatElem Type]
    -> ScatterSpec VName
    -> [(PatElem Type, (ShapeBase SubExp, Int, VName))]
    forall a b. [a] -> [b] -> [(a, b)]
    zip (Pat Type -> [PatElem Type]
    forall dec. Pat dec -> [PatElem dec]
    patElems Pat Type
    scat_pat) ScatterSpec VName
    scat_out
          wacc_stm <- [IotaInp]
    -> [RshpInp]
    -> [(PatElem (LetDec SOACS), (ShapeBase SubExp, Int, VName))]
    -> StmAux (ExpDec SOACS)
    -> Lambda SOACS
    -> m (Stm SOACS)
    forall (m :: * -> *).
    (HasScope SOACS m, MonadFreshNames m) =>
    [IotaInp]
    -> [RshpInp]
    -> [(PatElem (LetDec SOACS), (ShapeBase SubExp, Int, VName))]
    -> StmAux (ExpDec SOACS)
    -> Lambda SOACS
    -> m (Stm SOACS)
    mkWithAccStm [((VName, Param Type), (SubExp, SubExp, SubExp, IntType))]
    [IotaInp]
    rep_iotas [((VName, Param Type), (ShapeBase SubExp, ShapeBase SubExp, Type))]
    [RshpInp]
    rep_rshps [(PatElem Type, (ShapeBase SubExp, Int, VName))]
    [(PatElem (LetDec SOACS), (ShapeBase SubExp, Int, VName))]
    cons_patels_outs StmAux (ExpDec SOACS)
    scat_aux Lambda SOACS
    scat_lam
          let all_cert_rshp = [Certs] -> Certs
    forall a. Monoid a => [a] -> a
    mconcat [Certs]
    certs_rshps
              aux = Stm SOACS -> StmAux (ExpDec SOACS)
    forall rep. Stm rep -> StmAux (ExpDec rep)
    stmAux Stm SOACS
    wacc_stm
              aux' = StmAux ()
    aux {stmAuxCerts = all_cert_rshp <> stmAuxCerts aux}
              wacc_stm' = Stm SOACS
    wacc_stm {stmAux = aux'}
              -- get the input deps of iotas
              fiot a
    acc (a
    _, b
    _, c
    _, a
    inp_deps_iot) =
                a
    acc a -> a -> a
    forall a. Semigroup a => a -> a -> a
    <> a
    inp_deps_iot
              deps_of_iotas = (Adj EdgeT -> Context NodeT EdgeT -> Adj EdgeT)
    -> Adj EdgeT -> [Context NodeT EdgeT] -> Adj EdgeT
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl Adj EdgeT -> Context NodeT EdgeT -> Adj EdgeT
    forall {a} {a} {b} {c}. Semigroup a => a -> (a, b, c, a) -> a
    fiot Adj EdgeT
    forall a. Monoid a => a
    mempty [Context NodeT EdgeT]
    ctxs_iots
              --
              iota_nms = [VName] -> Names
    namesFromList ([VName] -> Names) -> [VName] -> Names
    forall a b. (a -> b) -> a -> b
    $ (((VName, Param Type), (SubExp, SubExp, SubExp, IntType)) -> VName)
    -> [((VName, Param Type), (SubExp, SubExp, SubExp, IntType))]
    -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map ((VName, Param Type) -> VName
    forall a b. (a, b) -> a
    fst ((VName, Param Type) -> VName)
    -> (((VName, Param Type), (SubExp, SubExp, SubExp, IntType))
        -> (VName, Param Type))
    -> ((VName, Param Type), (SubExp, SubExp, SubExp, IntType))
    -> VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ((VName, Param Type), (SubExp, SubExp, SubExp, IntType))
    -> (VName, Param Type)
    forall a b. (a, b) -> a
    fst) [((VName, Param Type), (SubExp, SubExp, SubExp, IntType))]
    rep_iotas
              inp_deps_wo_iotas = ((EdgeT, Int) -> Bool) -> Adj EdgeT -> Adj EdgeT
    forall a. (a -> Bool) -> [a] -> [a]
    filter ((VName -> Names -> Bool
    `notNameIn` Names
    iota_nms) (VName -> Bool) -> ((EdgeT, Int) -> VName) -> (EdgeT, Int) -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . EdgeT -> VName
    getName (EdgeT -> VName)
    -> ((EdgeT, Int) -> EdgeT) -> (EdgeT, Int) -> VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (EdgeT, Int) -> EdgeT
    forall a b. (a, b) -> a
    fst) Adj EdgeT
    inp_deps
              -- generate a new node for the with-acc-stmt and its associated context:
              --   add the inp-deps of iotas but remove the iota themselves from deps.
              new_withacc_nT = Stm SOACS -> NodeT
    StmNode Stm SOACS
    wacc_stm'
              inp_deps' = Adj EdgeT
    inp_deps_wo_iotas Adj EdgeT -> Adj EdgeT -> Adj EdgeT
    forall a. Semigroup a => a -> a -> a
    <> Adj EdgeT
    deps_of_iotas
              new_withacc_ctx = (Adj EdgeT
    out_deps, Int
    scat_node_id, NodeT
    new_withacc_nT, Adj EdgeT
    inp_deps')
              -- construct the new WithAcc node/graph; do we need to use `fusedSomething` ??
              new_node = Context NodeT EdgeT -> Int
    forall a b. Context a b -> Int
    G.node' Context NodeT EdgeT
    new_withacc_ctx
              dg' = DepGraph
    dg {dgGraph = new_withacc_ctx G.& G.delNodes [new_node] g}
          -- result
          pure $ Just dg'
      where
        --
        getNodeTfromCtx :: (a, b, c, d) -> c
    getNodeTfromCtx (a
    _, b
    _, c
    nT, d
    _) = c
    nT
        findCtxOf :: [(a, EdgeT)] -> VName -> Maybe (a, EdgeT)
    findCtxOf [(a, EdgeT)]
    ctxes VName
    nm
          | [(a, EdgeT)
    ctxe] <- ((a, EdgeT) -> Bool) -> [(a, EdgeT)] -> [(a, EdgeT)]
    forall a. (a -> Bool) -> [a] -> [a]
    filter (\(a, EdgeT)
    x -> VName
    nm VName -> VName -> Bool
    forall a. Eq a => a -> a -> Bool
    == EdgeT -> VName
    getName ((a, EdgeT) -> EdgeT
    forall a b. (a, b) -> b
    snd (a, EdgeT)
    x)) [(a, EdgeT)]
    ctxes =
              (a, EdgeT) -> Maybe (a, EdgeT)
    forall a. a -> Maybe a
    Just (a, EdgeT)
    ctxe
        findCtxOf [(a, EdgeT)]
    _ VName
    _ = Maybe (a, EdgeT)
    forall a. Maybe a
    Nothing
        pairUp :: [(DepContext, EdgeT)] -> H.Input -> Maybe (DepContext, (H.Input, NodeT))
        pairUp :: [(Context NodeT EdgeT, EdgeT)]
    -> Input -> Maybe (Context NodeT EdgeT, (Input, NodeT))
    pairUp [(Context NodeT EdgeT, EdgeT)]
    ctxes inp :: Input
    inp@(H.Input ArrayTransforms
    _arrtrsfs VName
    nm Type
    _tp)
          | Just (ctx :: Context NodeT EdgeT
    ctx@(Adj EdgeT
    _, Int
    _, NodeT
    nT, Adj EdgeT
    _), EdgeT
    _) <- [(Context NodeT EdgeT, EdgeT)]
    -> VName -> Maybe (Context NodeT EdgeT, EdgeT)
    forall {a}. [(a, EdgeT)] -> VName -> Maybe (a, EdgeT)
    findCtxOf [(Context NodeT EdgeT, EdgeT)]
    ctxes VName
    nm =
              (Context NodeT EdgeT, (Input, NodeT))
    -> Maybe (Context NodeT EdgeT, (Input, NodeT))
    forall a. a -> Maybe a
    Just (Context NodeT EdgeT
    ctx, (Input
    inp, NodeT
    nT))
        pairUp [(Context NodeT EdgeT, EdgeT)]
    _ Input
    _ = Maybe (Context NodeT EdgeT, (Input, NodeT))
    forall a. Maybe a
    Nothing
        --
        isIota :: NodeT -> Bool
        isIota :: NodeT -> Bool
    isIota (StmNode (Let Pat (LetDec SOACS)
    _ StmAux (ExpDec SOACS)
    _ (BasicOp (Iota {})))) = Bool
    True
        isIota NodeT
    _ = Bool
    False
        --
        getRepIota :: ((H.Input, NodeT), LParam SOACS) -> Maybe IotaInp
        getRepIota :: ((Input, NodeT), LParam SOACS) -> Maybe IotaInp
    getRepIota ((H.Input ArrayTransforms
    iottrsf VName
    arr_nm Type
    _arr_tp, NodeT
    nt), LParam SOACS
    farg)
          | ArrayTransforms
    forall a. Monoid a => a
    mempty ArrayTransforms -> ArrayTransforms -> Bool
    forall a. Eq a => a -> a -> Bool
    == ArrayTransforms
    iottrsf,
            StmNode (Let Pat (LetDec SOACS)
    _ StmAux (ExpDec SOACS)
    _ (BasicOp (Iota SubExp
    n SubExp
    x SubExp
    s IntType
    Int64))) <- NodeT
    nt =
              ((VName, Param Type), (SubExp, SubExp, SubExp, IntType))
    -> Maybe ((VName, Param Type), (SubExp, SubExp, SubExp, IntType))
    forall a. a -> Maybe a
    Just ((VName
    arr_nm, Param Type
    LParam SOACS
    farg), (SubExp
    n, SubExp
    x, SubExp
    s, IntType
    Int64))
        getRepIota ((Input, NodeT), LParam SOACS)
    _ = Maybe ((VName, Param Type), (SubExp, SubExp, SubExp, IntType))
    Maybe IotaInp
    forall a. Maybe a
    Nothing
        --
        getRepRshpArr :: ((H.Input, NodeT), LParam SOACS) -> Maybe (RshpInp, Certs)
        getRepRshpArr :: ((Input, NodeT), LParam SOACS) -> Maybe (RshpInp, Certs)
    getRepRshpArr ((H.Input ArrayTransforms
    outtrsf VName
    arr_nm Type
    arr_tp, NodeT
    _nt), LParam SOACS
    farg)
          | ArrayTransform
    rshp_trsfm H.:< ArrayTransforms
    other_trsfms <- ArrayTransforms -> ViewF
    H.viewf ArrayTransforms
    outtrsf,
            H.Reshape StmAux ()
    aux NewShape SubExp
    shp_flat <- ArrayTransform
    rshp_trsfm,
            ReshapeKind
    ReshapeArbitrary <- NewShape SubExp -> ReshapeKind
    reshapeKind NewShape SubExp
    shp_flat,
            ArrayTransforms
    other_trsfms ArrayTransforms -> ArrayTransforms -> Bool
    forall a. Eq a => a -> a -> Bool
    == ArrayTransforms
    forall a. Monoid a => a
    mempty,
            Type
    eltp <- Param Type -> Type
    forall dec. Param dec -> dec
    paramDec Param Type
    LParam SOACS
    farg,
            Just ShapeBase SubExp
    shp_flat' <- Type -> ShapeBase SubExp -> Maybe (ShapeBase SubExp)
    forall {b} {u}.
    Eq b =>
    TypeBase (ShapeBase b) u -> ShapeBase b -> Maybe (ShapeBase b)
    checkShp Type
    eltp (ShapeBase SubExp -> Maybe (ShapeBase SubExp))
    -> ShapeBase SubExp -> Maybe (ShapeBase SubExp)
    forall a b. (a -> b) -> a -> b
    $ NewShape SubExp -> ShapeBase SubExp
    forall d. NewShape d -> ShapeBase d
    newShape NewShape SubExp
    shp_flat,
            Array PrimType
    _ptp ShapeBase SubExp
    shp_unflat NoUniqueness
    _ <- Type
    arr_tp,
            Just ShapeBase SubExp
    shp_unflat' <- Type -> ShapeBase SubExp -> Maybe (ShapeBase SubExp)
    forall {b} {u}.
    Eq b =>
    TypeBase (ShapeBase b) u -> ShapeBase b -> Maybe (ShapeBase b)
    checkShp Type
    eltp ShapeBase SubExp
    shp_unflat,
            ShapeBase SubExp -> Int
    forall a. ArrayShape a => a -> Int
    shapeRank ShapeBase SubExp
    shp_flat' Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    1,
            ShapeBase SubExp -> Int
    forall a. ArrayShape a => a -> Int
    shapeRank ShapeBase SubExp
    shp_flat' Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    < ShapeBase SubExp -> Int
    forall a. ArrayShape a => a -> Int
    shapeRank ShapeBase SubExp
    shp_unflat' =
              (((VName, Param Type), (ShapeBase SubExp, ShapeBase SubExp, Type)),
     Certs)
    -> Maybe
         (((VName, Param Type), (ShapeBase SubExp, ShapeBase SubExp, Type)),
          Certs)
    forall a. a -> Maybe a
    Just (((VName
    arr_nm, Param Type
    LParam SOACS
    farg), (ShapeBase SubExp
    shp_flat', ShapeBase SubExp
    shp_unflat', Type
    eltp)), StmAux () -> Certs
    forall dec. StmAux dec -> Certs
    stmAuxCerts StmAux ()
    aux)
        getRepRshpArr ((Input, NodeT), LParam SOACS)
    _ = Maybe
      (((VName, Param Type), (ShapeBase SubExp, ShapeBase SubExp, Type)),
       Certs)
    Maybe (RshpInp, Certs)
    forall a. Maybe a
    Nothing
        --
        checkShp :: TypeBase (ShapeBase b) u -> ShapeBase b -> Maybe (ShapeBase b)
    checkShp (Prim PrimType
    _) ShapeBase b
    shp_arr = ShapeBase b -> Maybe (ShapeBase b)
    forall a. a -> Maybe a
    Just ShapeBase b
    shp_arr
        checkShp (Array PrimType
    _ptp ShapeBase b
    shp_elm u
    _) ShapeBase b
    shp_arr =
          let dims_elm :: [b]
    dims_elm = ShapeBase b -> [b]
    forall d. ShapeBase d -> [d]
    shapeDims ShapeBase b
    shp_elm
              dims_arr :: [b]
    dims_arr = ShapeBase b -> [b]
    forall d. ShapeBase d -> [d]
    shapeDims ShapeBase b
    shp_arr
              (Int
    m, Int
    n) = ([b] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [b]
    dims_elm, [b] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [b]
    dims_arr)
              shp' :: ShapeBase b
    shp' = [b] -> ShapeBase b
    forall d. [d] -> ShapeBase d
    Shape ([b] -> ShapeBase b) -> [b] -> ShapeBase b
    forall a b. (a -> b) -> a -> b
    $ Int -> [b] -> [b]
    forall a. Int -> [a] -> [a]
    take (Int
    n Int -> Int -> Int
    forall a. Num a => a -> a -> a
    - Int
    m) [b]
    dims_arr
              dims_com :: [b]
    dims_com = Int -> [b] -> [b]
    forall a. Int -> [a] -> [a]
    drop (Int
    n Int -> Int -> Int
    forall a. Num a => a -> a -> a
    - Int
    m) [b]
    dims_arr
           in if ((b, b) -> Bool) -> [(b, b)] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all ((b -> b -> Bool) -> (b, b) -> Bool
    forall a b c. (a -> b -> c) -> (a, b) -> c
    uncurry b -> b -> Bool
    forall a. Eq a => a -> a -> Bool
    (==)) ([b] -> [b] -> [(b, b)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [b]
    dims_com [b]
    dims_elm)
                then ShapeBase b -> Maybe (ShapeBase b)
    forall a. a -> Maybe a
    Just ShapeBase b
    shp'
                else Maybe (ShapeBase b)
    forall a. Maybe a
    Nothing
        checkShp TypeBase (ShapeBase b) u
    _ ShapeBase b
    _ = Maybe (ShapeBase b)
    forall a. Maybe a
    Nothing
    -- default fails:
    ruleMFScat DepNode
    _ DepGraph
    _ = Maybe DepGraph -> m (Maybe DepGraph)
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Maybe DepGraph
    forall a. Maybe a
    Nothing

    checkSafeAndProfitable :: DepGraph -> G.Node -> [DepContext] -> [DepContext] -> Bool
    checkSafeAndProfitable :: DepGraph
    -> Int -> [Context NodeT EdgeT] -> [Context NodeT EdgeT] -> Bool
    checkSafeAndProfitable DepGraph
    dg Int
    scat_node_id ctxs_rshp :: [Context NodeT EdgeT]
    ctxs_rshp@(Context NodeT EdgeT
    _ : [Context NodeT EdgeT]
    _) [Context NodeT EdgeT]
    ctxs_cons =
      let all_deps :: Adj EdgeT
    all_deps = (Context NodeT EdgeT -> Adj EdgeT)
    -> [Context NodeT EdgeT] -> Adj EdgeT
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap (\(Adj EdgeT
    x, Int
    _, NodeT
    _, Adj EdgeT
    _) -> Adj EdgeT
    x) ([Context NodeT EdgeT] -> Adj EdgeT)
    -> [Context NodeT EdgeT] -> Adj EdgeT
    forall a b. (a -> b) -> a -> b
    $ [Context NodeT EdgeT]
    ctxs_rshp [Context NodeT EdgeT]
    -> [Context NodeT EdgeT] -> [Context NodeT EdgeT]
    forall a. [a] -> [a] -> [a]
    ++ [Context NodeT EdgeT]
    ctxs_cons
          prof1 :: Bool
    prof1 = ((EdgeT, Int) -> Bool) -> Adj EdgeT -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all (\(EdgeT
    _, Int
    dep_id) -> Int
    dep_id Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    scat_node_id) Adj EdgeT
    all_deps
          -- \^ scatter is the sole target to all consume & unflatten-reshape deps
          (Adj EdgeT
    _, Int
    map_node_id, NodeT
    map_nT, Adj EdgeT
    _) = [Context NodeT EdgeT] -> Context NodeT EdgeT
    forall a. HasCallStack => [a] -> a
    head [Context NodeT EdgeT]
    ctxs_rshp
          prof2 :: Bool
    prof2 = (Context NodeT EdgeT -> Bool) -> [Context NodeT EdgeT] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all (\(Adj EdgeT
    _, Int
    nid, NodeT
    _, Adj EdgeT
    _) -> Int
    nid Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    map_node_id) [Context NodeT EdgeT]
    ctxs_rshp
          prof3 :: Bool
    prof3 = NodeT -> Bool
    isMap NodeT
    map_nT
          -- \^ all reshapes come from the same node, which is a map
          safe :: Bool
    safe = DepGraph -> Int -> Int -> Bool
    vFusionFeasability DepGraph
    dg Int
    map_node_id Int
    scat_node_id
       in Bool
    safe Bool -> Bool -> Bool
    && Bool
    prof1 Bool -> Bool -> Bool
    && Bool
    prof2 Bool -> Bool -> Bool
    && Bool
    prof3
      where
        isMap :: NodeT -> Bool
    isMap NodeT
    nT
          | SoacNode ArrayTransforms
    out_trsfs Pat Type
    _pat SOAC SOACS
    soac StmAux (ExpDec SOACS)
    _ <- NodeT
    nT,
            H.Screma SubExp
    _ [Input]
    _ ScremaForm SOACS
    form <- SOAC SOACS
    soac,
            ScremaForm Lambda SOACS
    _ [] [] <- ScremaForm SOACS
    form =
              ArrayTransforms -> Bool
    H.nullTransforms ArrayTransforms
    out_trsfs
        isMap NodeT
    _ = Bool
    False
    checkSafeAndProfitable DepGraph
    _ Int
    _ [Context NodeT EdgeT]
    _ [Context NodeT EdgeT]
    _ = Bool
    False

    -- | produces the withAcc statement that constitutes the translation of
    --   the scater o flatten o map composition in which the map inputs are
    --   reshaped in the same way
    mkWithAccStm ::
      (HasScope SOACS m, MonadFreshNames m) =>
      [IotaInp] ->
      [RshpInp] ->
      [(PatElem (LetDec SOACS), (Shape, Int, VName))] ->
      StmAux (ExpDec SOACS) ->
      Lambda SOACS ->
      m (Stm SOACS)
    mkWithAccStm :: forall (m :: * -> *).
    (HasScope SOACS m, MonadFreshNames m) =>
    [IotaInp]
    -> [RshpInp]
    -> [(PatElem (LetDec SOACS), (ShapeBase SubExp, Int, VName))]
    -> StmAux (ExpDec SOACS)
    -> Lambda SOACS
    -> m (Stm SOACS)
    mkWithAccStm [IotaInp]
    iota_inps [RshpInp]
    rshp_inps [(PatElem (LetDec SOACS), (ShapeBase SubExp, Int, VName))]
    cons_patels_outs StmAux (ExpDec SOACS)
    scatter_aux Lambda SOACS
    scatter_lam
      -- iotas are assumed to operate on Int64 values
      -- ToDo: maybe simplify rshp_inps
      --       check that the unflat shape is the same across reshapes
      --       check that the rank of the unflatten shape is higher than the flatten
      | RshpInp
    rshp_inp : [RshpInp]
    _ <- [RshpInp]
    rshp_inps,
        ((VName, LParam SOACS)
    _, (ShapeBase SubExp
    _, ShapeBase SubExp
    s_unflat, Type
    _)) <- RshpInp
    rshp_inp,
        (SubExp
    _ : [SubExp]
    _) <- ShapeBase SubExp -> [SubExp]
    forall d. ShapeBase d -> [d]
    shapeDims ShapeBase SubExp
    s_unflat = do
          --
          (cert_params, acc_params) <- ([(Param Type, Param Type)] -> ([Param Type], [Param Type]))
    -> m [(Param Type, Param Type)] -> m ([Param Type], [Param Type])
    forall a b. (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap [(Param Type, Param Type)] -> ([Param Type], [Param Type])
    forall a b. [(a, b)] -> ([a], [b])
    unzip (m [(Param Type, Param Type)] -> m ([Param Type], [Param Type]))
    -> m [(Param Type, Param Type)] -> m ([Param Type], [Param Type])
    forall a b. (a -> b) -> a -> b
    $
            [(PatElem Type, (ShapeBase SubExp, Int, VName))]
    -> ((PatElem Type, (ShapeBase SubExp, Int, VName))
        -> m (Param Type, Param Type))
    -> m [(Param Type, Param Type)]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM [(PatElem Type, (ShapeBase SubExp, Int, VName))]
    [(PatElem (LetDec SOACS), (ShapeBase SubExp, Int, VName))]
    cons_patels_outs (((PatElem Type, (ShapeBase SubExp, Int, VName))
      -> m (Param Type, Param Type))
     -> m [(Param Type, Param Type)])
    -> ((PatElem Type, (ShapeBase SubExp, Int, VName))
        -> m (Param Type, Param Type))
    -> m [(Param Type, Param Type)]
    forall a b. (a -> b) -> a -> b
    $ \(PatElem Type
    patel, (ShapeBase SubExp
    shp, Int
    _, VName
    nm)) -> do
              cert_param <- String -> Type -> m (Param Type)
    forall (m :: * -> *) dec.
    MonadFreshNames m =>
    String -> dec -> m (Param dec)
    newParam String
    "acc_cert_p" (Type -> m (Param Type)) -> Type -> m (Param Type)
    forall a b. (a -> b) -> a -> b
    $ PrimType -> Type
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    Unit
              let arr_tp = PatElem Type -> Type
    forall dec. Typed dec => PatElem dec -> Type
    patElemType PatElem Type
    patel
                  acc_tp = Int -> Type -> Type
    forall u.
    Int
    -> TypeBase (ShapeBase SubExp) u -> TypeBase (ShapeBase SubExp) u
    stripArray (ShapeBase SubExp -> Int
    forall a. ArrayShape a => a -> Int
    shapeRank ShapeBase SubExp
    shp) Type
    arr_tp
              acc_param <-
                newParam (baseString nm) $
                  Acc (paramName cert_param) shp [acc_tp] NoUniqueness
              pure (cert_param, acc_param)
          let cons_params_outs = [Param Type]
    -> ScatterSpec VName
    -> [(Param Type, (ShapeBase SubExp, Int, VName))]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Param Type]
    acc_params (ScatterSpec VName
     -> [(Param Type, (ShapeBase SubExp, Int, VName))])
    -> ScatterSpec VName
    -> [(Param Type, (ShapeBase SubExp, Int, VName))]
    forall a b. (a -> b) -> a -> b
    $ ((PatElem Type, (ShapeBase SubExp, Int, VName))
     -> (ShapeBase SubExp, Int, VName))
    -> [(PatElem Type, (ShapeBase SubExp, Int, VName))]
    -> ScatterSpec VName
    forall a b. (a -> b) -> [a] -> [b]
    map (PatElem Type, (ShapeBase SubExp, Int, VName))
    -> (ShapeBase SubExp, Int, VName)
    forall a b. (a, b) -> b
    snd [(PatElem Type, (ShapeBase SubExp, Int, VName))]
    [(PatElem (LetDec SOACS), (ShapeBase SubExp, Int, VName))]
    cons_patels_outs
          acc_bdy <- mkWithAccBdy s_unflat iota_inps rshp_inps cons_params_outs scatter_lam
          let withacc_lam =
                Lambda
                  { lambdaParams :: [LParam SOACS]
    lambdaParams = [Param Type]
    cert_params [Param Type] -> [Param Type] -> [Param Type]
    forall a. [a] -> [a] -> [a]
    ++ [Param Type]
    acc_params,
                    lambdaReturnType :: [Type]
    lambdaReturnType = (Param Type -> Type) -> [Param Type] -> [Type]
    forall a b. (a -> b) -> [a] -> [b]
    map Param Type -> Type
    forall dec. Param dec -> dec
    paramDec [Param Type]
    acc_params,
                    lambdaBody :: Body SOACS
    lambdaBody = Body SOACS
    acc_bdy
                  }
              withacc_inps = ((PatElem Type, (ShapeBase SubExp, Int, VName))
     -> (ShapeBase SubExp, [VName], Maybe a))
    -> [(PatElem Type, (ShapeBase SubExp, Int, VName))]
    -> [(ShapeBase SubExp, [VName], Maybe a)]
    forall a b. (a -> b) -> [a] -> [b]
    map (\(PatElem Type
    _, (ShapeBase SubExp
    shp, Int
    _, VName
    nm)) -> (ShapeBase SubExp
    shp, [VName
    nm], Maybe a
    forall a. Maybe a
    Nothing)) [(PatElem Type, (ShapeBase SubExp, Int, VName))]
    [(PatElem (LetDec SOACS), (ShapeBase SubExp, Int, VName))]
    cons_patels_outs
              withacc_pat = [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
    $ ((PatElem Type, (ShapeBase SubExp, Int, VName)) -> PatElem Type)
    -> [(PatElem Type, (ShapeBase SubExp, Int, VName))]
    -> [PatElem Type]
    forall a b. (a -> b) -> [a] -> [b]
    map (PatElem Type, (ShapeBase SubExp, Int, VName)) -> PatElem Type
    forall a b. (a, b) -> a
    fst [(PatElem Type, (ShapeBase SubExp, Int, VName))]
    [(PatElem (LetDec SOACS), (ShapeBase SubExp, Int, VName))]
    cons_patels_outs
              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)
    withacc_pat StmAux (ExpDec SOACS)
    scatter_aux (Exp SOACS -> Stm SOACS) -> Exp SOACS -> Stm SOACS
    forall a b. (a -> b) -> a -> b
    $
                  [WithAccInput SOACS] -> Lambda SOACS -> Exp SOACS
    forall rep. [WithAccInput rep] -> Lambda rep -> Exp rep
    WithAcc [WithAccInput SOACS]
    forall {a}. [(ShapeBase SubExp, [VName], Maybe a)]
    withacc_inps Lambda SOACS
    withacc_lam
          pure stm
    mkWithAccStm [IotaInp]
    _ [RshpInp]
    _ [(PatElem (LetDec SOACS), (ShapeBase SubExp, Int, VName))]
    _ StmAux (ExpDec SOACS)
    _ Lambda SOACS
    _ =
      String -> m (Stm SOACS)
    forall a. HasCallStack => String -> a
    error String
    "Unreachable case reached!"

    -- | Wrapper function for constructing the body of the withAcc
    --   translation of the scatter
    mkWithAccBdy ::
      (HasScope SOACS m, MonadFreshNames m) =>
      Shape ->
      [IotaInp] ->
      [RshpInp] ->
      [(LParam SOACS, (Shape, Int, VName))] ->
      Lambda SOACS ->
      m (Body SOACS)
    mkWithAccBdy :: forall (m :: * -> *).
    (HasScope SOACS m, MonadFreshNames m) =>
    ShapeBase SubExp
    -> [IotaInp]
    -> [RshpInp]
    -> [(LParam SOACS, (ShapeBase SubExp, Int, VName))]
    -> Lambda SOACS
    -> m (Body SOACS)
    mkWithAccBdy ShapeBase SubExp
    shp [IotaInp]
    iota_inps [RshpInp]
    rshp_inps [(LParam SOACS, (ShapeBase SubExp, Int, VName))]
    cons_params_outs Lambda SOACS
    scat_lam = do
      let cons_ps :: [Param Type]
    cons_ps = ((Param Type, (ShapeBase SubExp, Int, VName)) -> Param Type)
    -> [(Param Type, (ShapeBase SubExp, Int, VName))] -> [Param Type]
    forall a b. (a -> b) -> [a] -> [b]
    map (Param Type, (ShapeBase SubExp, Int, VName)) -> Param Type
    forall a b. (a, b) -> a
    fst [(Param Type, (ShapeBase SubExp, Int, VName))]
    [(LParam SOACS, (ShapeBase SubExp, Int, VName))]
    cons_params_outs
          scat_res_info :: ScatterSpec VName
    scat_res_info = ((Param Type, (ShapeBase SubExp, Int, VName))
     -> (ShapeBase SubExp, Int, VName))
    -> [(Param Type, (ShapeBase SubExp, Int, VName))]
    -> ScatterSpec VName
    forall a b. (a -> b) -> [a] -> [b]
    map (Param Type, (ShapeBase SubExp, Int, VName))
    -> (ShapeBase SubExp, Int, VName)
    forall a b. (a, b) -> b
    snd [(Param Type, (ShapeBase SubExp, Int, VName))]
    [(LParam SOACS, (ShapeBase SubExp, Int, VName))]
    cons_params_outs
          static_arg :: ([((VName, Param Type), (SubExp, SubExp, SubExp, IntType))],
     [((VName, Param Type),
       (ShapeBase SubExp, ShapeBase SubExp, Type))],
     ScatterSpec VName, Lambda SOACS)
    static_arg = ([((VName, Param Type), (SubExp, SubExp, SubExp, IntType))]
    [IotaInp]
    iota_inps, [((VName, Param Type), (ShapeBase SubExp, ShapeBase SubExp, Type))]
    [RshpInp]
    rshp_inps, ScatterSpec VName
    scat_res_info, Lambda SOACS
    scat_lam)
          mkParam :: ((VName, b), (a, ShapeBase SubExp, Type)) -> Param Type
    mkParam ((VName
    nm, b
    _), (a
    _, ShapeBase SubExp
    s, Type
    t)) = Attrs -> VName -> Type -> Param Type
    forall dec. Attrs -> VName -> dec -> Param dec
    Param Attrs
    forall a. Monoid a => a
    mempty VName
    nm (Type -> ShapeBase SubExp -> Type
    arrayOfShape Type
    t ShapeBase SubExp
    s)
          rshp_ps :: [Param Type]
    rshp_ps = (((VName, Param Type), (ShapeBase SubExp, ShapeBase SubExp, Type))
     -> Param Type)
    -> [((VName, Param Type),
         (ShapeBase SubExp, ShapeBase SubExp, Type))]
    -> [Param Type]
    forall a b. (a -> b) -> [a] -> [b]
    map ((VName, Param Type), (ShapeBase SubExp, ShapeBase SubExp, Type))
    -> Param Type
    forall {b} {a}.
    ((VName, b), (a, ShapeBase SubExp, Type)) -> Param Type
    mkParam [((VName, Param Type), (ShapeBase SubExp, ShapeBase SubExp, Type))]
    [RshpInp]
    rshp_inps
      ([IotaInp], [RshpInp], ScatterSpec VName, Lambda SOACS)
    -> [SubExp]
    -> [SubExp]
    -> [VName]
    -> [LParam SOACS]
    -> [LParam SOACS]
    -> m (Body SOACS)
    forall (m :: * -> *).
    (HasScope SOACS m, MonadFreshNames m) =>
    ([IotaInp], [RshpInp], ScatterSpec VName, Lambda SOACS)
    -> [SubExp]
    -> [SubExp]
    -> [VName]
    -> [LParam SOACS]
    -> [LParam SOACS]
    -> m (Body SOACS)
    mkWithAccBdy' ([((VName, Param Type), (SubExp, SubExp, SubExp, IntType))],
     [((VName, Param Type),
       (ShapeBase SubExp, ShapeBase SubExp, Type))],
     ScatterSpec VName, Lambda SOACS)
    ([IotaInp], [RshpInp], ScatterSpec VName, Lambda SOACS)
    static_arg (ShapeBase SubExp -> [SubExp]
    forall d. ShapeBase d -> [d]
    shapeDims ShapeBase SubExp
    shp) [] [] [Param Type]
    [LParam SOACS]
    rshp_ps [Param Type]
    [LParam SOACS]
    cons_ps

    -- | builds a body that essentially consists of a map-nest with accumulators,
    --   i.e., one level for each level of the unflatten shape of scatter's reshaped
    --   input arrays
    mkWithAccBdy' ::
      (HasScope SOACS m, MonadFreshNames m) =>
      ([IotaInp], [RshpInp], [(Shape, Int, VName)], Lambda SOACS) ->
      [SubExp] ->
      [SubExp] ->
      [VName] ->
      [LParam SOACS] ->
      [LParam SOACS] ->
      m (Body SOACS)

    -- | the base case below addapts the scatter's lambda
    mkWithAccBdy' :: forall (m :: * -> *).
    (HasScope SOACS m, MonadFreshNames m) =>
    ([IotaInp], [RshpInp], ScatterSpec VName, Lambda SOACS)
    -> [SubExp]
    -> [SubExp]
    -> [VName]
    -> [LParam SOACS]
    -> [LParam SOACS]
    -> m (Body SOACS)
    mkWithAccBdy' ([IotaInp], [RshpInp], ScatterSpec VName, Lambda SOACS)
    static_arg [] [SubExp]
    dims_rev [VName]
    iot_par_nms [LParam SOACS]
    rshp_ps [LParam SOACS]
    cons_ps = do
      let ([IotaInp]
    iota_inps, [RshpInp]
    rshp_inps, ScatterSpec VName
    scat_res_info, Lambda SOACS
    scat_lam) = ([IotaInp], [RshpInp], ScatterSpec VName, Lambda SOACS)
    static_arg
          tp_int :: TypeBase shape u
    tp_int = PrimType -> TypeBase shape u
    forall shape u. PrimType -> TypeBase shape u
    Prim (PrimType -> TypeBase shape u) -> PrimType -> TypeBase shape u
    forall a b. (a -> b) -> a -> b
    $ IntType -> PrimType
    IntType IntType
    Int64
      scope <- m (Scope SOACS)
    forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
    askScope
      runBodyBuilder $ localScope (scope <> scopeOfLParams (rshp_ps ++ cons_ps)) $ do
        -- handle iota args
        let strides_rev = (TPrimExp Int64 VName
     -> TPrimExp Int64 VName -> TPrimExp Int64 VName)
    -> TPrimExp Int64 VName
    -> [TPrimExp Int64 VName]
    -> [TPrimExp Int64 VName]
    forall b a. (b -> a -> b) -> b -> [a] -> [b]
    scanl TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    (*) (SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    se1) ([TPrimExp Int64 VName] -> [TPrimExp Int64 VName])
    -> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
    forall a b. (a -> b) -> a -> b
    $ (SubExp -> TPrimExp Int64 VName)
    -> [SubExp] -> [TPrimExp Int64 VName]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExp -> TPrimExp Int64 VName
    pe64 [SubExp]
    dims_rev
            strides = [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
    forall a. HasCallStack => [a] -> [a]
    tail ([TPrimExp Int64 VName] -> [TPrimExp Int64 VName])
    -> [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
    forall a b. (a -> b) -> a -> b
    $ [TPrimExp Int64 VName] -> [TPrimExp Int64 VName]
    forall a. [a] -> [a]
    reverse [TPrimExp Int64 VName]
    strides_rev
            prods = (TPrimExp Int64 VName
     -> TPrimExp Int64 VName -> TPrimExp Int64 VName)
    -> [TPrimExp Int64 VName]
    -> [TPrimExp Int64 VName]
    -> [TPrimExp Int64 VName]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    (*) ((VName -> TPrimExp Int64 VName)
    -> [VName] -> [TPrimExp Int64 VName]
    forall a b. (a -> b) -> [a] -> [b]
    map VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 [VName]
    iot_par_nms) [TPrimExp Int64 VName]
    strides
            i_pe = [TPrimExp Int64 VName] -> TPrimExp Int64 VName
    forall a. Num a => [a] -> a
    forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
    sum [TPrimExp Int64 VName]
    prods
        i_norm <- letExp "iota_norm_arg" =<< toExp i_pe
        forM_ iota_inps $ \((VName, Param Type), (SubExp, SubExp, SubExp, IntType))
    arg -> do
          let ((VName
    _, Param Type
    i_par), (SubExp
    _, SubExp
    b, SubExp
    s, IntType
    _)) = ((VName, Param Type), (SubExp, SubExp, SubExp, IntType))
    arg
          i_new <- String
    -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    -> BuilderT SOACS (State VNameSource) VName
    forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m VName
    letExp String
    "tmp" (Exp SOACS -> BuilderT SOACS (State VNameSource) VName)
    -> BuilderT SOACS (State VNameSource) (Exp SOACS)
    -> BuilderT SOACS (State VNameSource) VName
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< TPrimExp Int64 VName
    -> BuilderT
         SOACS
         (State VNameSource)
         (Exp (Rep (BuilderT SOACS (State VNameSource))))
    forall a (m :: * -> *).
    (ToExp a, MonadBuilder m) =>
    a -> m (Exp (Rep m))
    forall (m :: * -> *).
    MonadBuilder m =>
    TPrimExp Int64 VName -> m (Exp (Rep m))
    toExp (SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    b TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    + VName -> TPrimExp Int64 VName
    forall a. a -> TPrimExp Int64 a
    le64 VName
    i_norm TPrimExp Int64 VName
    -> TPrimExp Int64 VName -> TPrimExp Int64 VName
    forall a. Num a => a -> a -> a
    * SubExp -> TPrimExp Int64 VName
    pe64 SubExp
    s)
          letBind (Pat [PatElem (paramName i_par) tp_int]) $ BasicOp $ SubExp $ Var i_new
        -- handle rshp args
        let rshp_lam_args = (((VName, Param Type), (ShapeBase SubExp, ShapeBase SubExp, Type))
     -> Param Type)
    -> [((VName, Param Type),
         (ShapeBase SubExp, ShapeBase SubExp, Type))]
    -> [Param Type]
    forall a b. (a -> b) -> [a] -> [b]
    map ((VName, Param Type) -> Param Type
    forall a b. (a, b) -> b
    snd ((VName, Param Type) -> Param Type)
    -> (((VName, Param Type),
         (ShapeBase SubExp, ShapeBase SubExp, Type))
        -> (VName, Param Type))
    -> ((VName, Param Type),
        (ShapeBase SubExp, ShapeBase SubExp, Type))
    -> Param Type
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ((VName, Param Type), (ShapeBase SubExp, ShapeBase SubExp, Type))
    -> (VName, Param Type)
    forall a b. (a, b) -> a
    fst) [((VName, Param Type), (ShapeBase SubExp, ShapeBase SubExp, Type))]
    rshp_inps
        forM_ (zip rshp_lam_args rshp_ps) $ \(Param Type
    old_par, Param Type
    new_par) -> do
          let pat :: Pat Type
    pat = [PatElem Type] -> Pat Type
    forall dec. [PatElem dec] -> Pat dec
    Pat [VName -> Type -> PatElem Type
    forall dec. VName -> dec -> PatElem dec
    PatElem (Param Type -> VName
    forall dec. Param dec -> VName
    paramName Param Type
    old_par) (Param Type -> Type
    forall dec. Param dec -> dec
    paramDec Param Type
    old_par)]
          Pat (LetDec (Rep (BuilderT SOACS (State VNameSource))))
    -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    -> BuilderT SOACS (State VNameSource) ()
    forall (m :: * -> *).
    MonadBuilder m =>
    Pat (LetDec (Rep m)) -> Exp (Rep m) -> m ()
    letBind Pat Type
    Pat (LetDec (Rep (BuilderT SOACS (State VNameSource))))
    pat (Exp (Rep (BuilderT SOACS (State VNameSource)))
     -> BuilderT SOACS (State VNameSource) ())
    -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    -> BuilderT SOACS (State VNameSource) ()
    forall a b. (a -> b) -> a -> b
    $ BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource))))
    -> BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    forall a b. (a -> b) -> a -> b
    $ 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
    $ Param Type -> VName
    forall dec. Param dec -> VName
    paramName Param Type
    new_par
        -- add the body of the scatter's lambda
        mapM_ addStm $ bodyStms $ lambdaBody scat_lam
        -- add the withAcc update statements
        let iv_ses = ScatterSpec VName -> Result -> [(Result, SubExpRes)]
    forall array a.
    [(ShapeBase SubExp, Int, array)] -> [a] -> [([a], a)]
    groupScatterResults' ScatterSpec VName
    scat_res_info (Result -> [(Result, SubExpRes)])
    -> Result -> [(Result, SubExpRes)]
    forall a b. (a -> b) -> a -> b
    $ Body SOACS -> Result
    forall rep. Body rep -> Result
    bodyResult (Body SOACS -> Result) -> Body SOACS -> Result
    forall a b. (a -> b) -> a -> b
    $ Lambda SOACS -> Body SOACS
    forall rep. Lambda rep -> Body rep
    lambdaBody Lambda SOACS
    scat_lam
        res_nms <-
          forM (zip cons_ps iv_ses) $ \(Param Type
    cons_p, (Result
    i_ses, SubExpRes
    v_se)) -> do
            -- i_ses is a list
            let f :: VName -> SubExpRes -> m VName
    f VName
    nm_in SubExpRes
    i_se =
                  String -> Exp (Rep m) -> m VName
    forall (m :: * -> *).
    MonadBuilder m =>
    String -> Exp (Rep m) -> m VName
    letExp (VName -> String
    baseString VName
    nm_in) (Exp (Rep m) -> m VName) -> Exp (Rep m) -> m VName
    forall a b. (a -> b) -> a -> b
    $ BasicOp -> Exp (Rep m)
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp -> Exp (Rep m)) -> BasicOp -> Exp (Rep m)
    forall a b. (a -> b) -> a -> b
    $ Safety -> VName -> [SubExp] -> [SubExp] -> BasicOp
    UpdateAcc Safety
    Safe VName
    nm_in [SubExpRes -> SubExp
    resSubExp SubExpRes
    i_se] [SubExpRes -> SubExp
    resSubExp SubExpRes
    v_se]
            (VName -> SubExpRes -> BuilderT SOACS (State VNameSource) VName)
    -> VName -> Result -> BuilderT SOACS (State VNameSource) VName
    forall (t :: * -> *) (m :: * -> *) b a.
    (Foldable t, Monad m) =>
    (b -> a -> m b) -> b -> t a -> m b
    foldM VName -> SubExpRes -> BuilderT SOACS (State VNameSource) VName
    forall {m :: * -> *}.
    MonadBuilder m =>
    VName -> SubExpRes -> m VName
    f (Param Type -> VName
    forall dec. Param dec -> VName
    paramName Param Type
    cons_p) Result
    i_ses
        let lam_certs = (SubExpRes -> Certs) -> Result -> Certs
    forall m a. Monoid m => (a -> m) -> [a] -> m
    forall (t :: * -> *) m a.
    (Foldable t, Monoid m) =>
    (a -> m) -> t a -> m
    foldMap SubExpRes -> Certs
    resCerts (Result -> Certs) -> Result -> Certs
    forall a b. (a -> b) -> a -> b
    $ Body SOACS -> Result
    forall rep. Body rep -> Result
    bodyResult (Body SOACS -> Result) -> Body SOACS -> Result
    forall a b. (a -> b) -> a -> b
    $ Lambda SOACS -> Body SOACS
    forall rep. Lambda rep -> Body rep
    lambdaBody Lambda SOACS
    scat_lam
        pure $ map (SubExpRes lam_certs . Var) res_nms
    -- \| the recursive case builds a call to a map soac.
    mkWithAccBdy' ([IotaInp], [RshpInp], ScatterSpec VName, Lambda SOACS)
    static_arg (SubExp
    dim : [SubExp]
    dims) [SubExp]
    dims_rev [VName]
    iot_par_nms [LParam SOACS]
    rshp_ps [LParam SOACS]
    cons_ps = do
      scope <- m (Scope SOACS)
    forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
    askScope
      runBodyBuilder $ localScope (scope <> scopeOfLParams (rshp_ps ++ cons_ps)) $ do
        iota_arr <- letExp "iota_arr" $ BasicOp $ Iota dim se0 se1 Int64
        iota_p <- newParam "iota_arg" $ Prim $ IntType Int64
        rshp_ps' <- forM (zip [0 .. length rshp_ps - 1] (map paramDec rshp_ps)) $
          \(Int
    i, Type
    arr_tp) ->
            String -> Type -> BuilderT SOACS (State VNameSource) (Param Type)
    forall (m :: * -> *) dec.
    MonadFreshNames m =>
    String -> dec -> m (Param dec)
    newParam (String
    "rshp_arg_" String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ Int -> String
    forall a. Show a => a -> String
    show Int
    i) (Type -> BuilderT SOACS (State VNameSource) (Param Type))
    -> Type -> BuilderT SOACS (State VNameSource) (Param Type)
    forall a b. (a -> b) -> a -> b
    $ Int -> Type -> Type
    forall u.
    Int
    -> TypeBase (ShapeBase SubExp) u -> TypeBase (ShapeBase SubExp) u
    stripArray Int
    1 Type
    arr_tp
        cons_ps' <- forM (zip [0 .. length cons_ps - 1] (map paramDec cons_ps)) $
          \(Int
    i, Type
    arr_tp) ->
            String -> Type -> BuilderT SOACS (State VNameSource) (Param Type)
    forall (m :: * -> *) dec.
    MonadFreshNames m =>
    String -> dec -> m (Param dec)
    newParam (String
    "acc_arg_" String -> String -> String
    forall a. [a] -> [a] -> [a]
    ++ Int -> String
    forall a. Show a => a -> String
    show Int
    i) Type
    arr_tp
        map_lam_bdy <-
          mkWithAccBdy' static_arg dims (dim : dims_rev) (iot_par_nms ++ [paramName iota_p]) rshp_ps' cons_ps'
        let map_lam = [LParam SOACS] -> [Type] -> Body SOACS -> Lambda SOACS
    forall rep. [LParam rep] -> [Type] -> Body rep -> Lambda rep
    Lambda ([Param Type]
    rshp_ps' [Param Type] -> [Param Type] -> [Param Type]
    forall a. [a] -> [a] -> [a]
    ++ [Param Type
    iota_p] [Param Type] -> [Param Type] -> [Param Type]
    forall a. [a] -> [a] -> [a]
    ++ [Param Type]
    cons_ps') ((Param Type -> Type) -> [Param Type] -> [Type]
    forall a b. (a -> b) -> [a] -> [b]
    map Param Type -> Type
    forall dec. Param dec -> dec
    paramDec [Param Type]
    cons_ps') Body SOACS
    map_lam_bdy
            map_inps = (Param Type -> VName) -> [Param Type] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map Param Type -> VName
    forall dec. Param dec -> VName
    paramName [Param Type]
    [LParam SOACS]
    rshp_ps [VName] -> [VName] -> [VName]
    forall a. [a] -> [a] -> [a]
    ++ [VName
    iota_arr] [VName] -> [VName] -> [VName]
    forall a. [a] -> [a] -> [a]
    ++ (Param Type -> VName) -> [Param Type] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map Param Type -> VName
    forall dec. Param dec -> VName
    paramName [Param Type]
    [LParam SOACS]
    cons_ps
            map_soac = SubExp -> [VName] -> ScremaForm SOACS -> SOAC SOACS
    forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
    F.Screma SubExp
    dim [VName]
    map_inps (ScremaForm SOACS -> SOAC SOACS) -> ScremaForm SOACS -> SOAC SOACS
    forall a b. (a -> b) -> a -> b
    $ Lambda SOACS -> [Scan SOACS] -> [Reduce SOACS] -> ScremaForm SOACS
    forall rep.
    Lambda rep -> [Scan rep] -> [Reduce rep] -> ScremaForm rep
    ScremaForm Lambda SOACS
    map_lam [] []
        res_nms <- letTupExp "acc_res" $ Op map_soac
        pure $ map (subExpRes . Var) res_nms

    ---------------------------------------------------
    --- II. WithAcc-WithAcc Fusion
    ---------------------------------------------------

    -- | Local helper type that tuples together:
    --   1.   the pattern element corresponding to one withacc input
    --   2.   the withacc input
    --   3-5  withacc's lambda corresponding acc-certificate param,
    --           argument param and result name
    type AccTup =
      ( [PatElem (LetDec SOACS)],
        WithAccInput SOACS,
        LParam SOACS,
        LParam SOACS,
        (VName, Certs)
      )

    accTup1 :: AccTup -> [PatElem (LetDec SOACS)]
    accTup1 :: AccTup -> [PatElem (LetDec SOACS)]
    accTup1 ([PatElem (LetDec SOACS)]
    a, WithAccInput SOACS
    _, LParam SOACS
    _, LParam SOACS
    _, (VName, Certs)
    _) = [PatElem (LetDec SOACS)]
    a

    accTup2 :: AccTup -> WithAccInput SOACS
    accTup2 :: AccTup -> WithAccInput SOACS
    accTup2 ([PatElem (LetDec SOACS)]
    _, WithAccInput SOACS
    a, LParam SOACS
    _, LParam SOACS
    _, (VName, Certs)
    _) = WithAccInput SOACS
    a

    accTup3 :: AccTup -> LParam SOACS
    accTup3 :: AccTup -> LParam SOACS
    accTup3 ([PatElem (LetDec SOACS)]
    _, WithAccInput SOACS
    _, LParam SOACS
    a, LParam SOACS
    _, (VName, Certs)
    _) = LParam SOACS
    a

    accTup4 :: AccTup -> LParam SOACS
    accTup4 :: AccTup -> LParam SOACS
    accTup4 ([PatElem (LetDec SOACS)]
    _, WithAccInput SOACS
    _, LParam SOACS
    _, LParam SOACS
    a, (VName, Certs)
    _) = LParam SOACS
    a

    accTup5 :: AccTup -> (VName, Certs)
    accTup5 :: AccTup -> (VName, Certs)
    accTup5 ([PatElem (LetDec SOACS)]
    _, WithAccInput SOACS
    _, LParam SOACS
    _, LParam SOACS
    _, (VName, Certs)
    a) = (VName, Certs)
    a

    -- | Simple case for fusing two withAccs (can be extended):
    --    let (b1, ..., bm, x1, ..., xq) = withAcc a1 ... am lam1
    --    let (d1, ..., dn, y1, ..., yp) = withAcc c1 ... cn lam2
    -- Notation: `b1 ... bm` are the accumulator results of the
    --     first withAcc and `d1, ..., dn` of the second withAcc.
    --     `x1 ... xq` and `y1, ..., yp` are non-accumulator results.
    -- Conservative conditions:
    --   1. for any bi (i=1..m) either `bi IN {c1, ..., cm}` OR
    --        `bi NOT-IN FV(lam2)`, i.e., perfect producer-consumer
    --        relation on accums. Of course the binary-op should
    --        be the same.
    --   2. The `bs` that are also accumulated upon in lam2
    --        do NOT belong to the `infusible` set (they are destroyed)
    --   3. x1 ... xq do not overlap with c1 ... cn
    -- Fusion will create one withacc that accumulates on the
    --   union of `a1 ... am` and `c1 ... cn` and returns, in addition
    --   to the accumulator arrays the union of regular variables
    --   `x1 ... xq` and `y1, ..., yp`
    tryFuseWithAccs ::
      (HasScope SOACS m, MonadFreshNames m) =>
      [VName] ->
      Stm SOACS ->
      Stm SOACS ->
      Maybe (m (Stm SOACS))
    tryFuseWithAccs :: forall (m :: * -> *).
    (HasScope SOACS m, MonadFreshNames m) =>
    [VName] -> Stm SOACS -> Stm SOACS -> Maybe (m (Stm SOACS))
    tryFuseWithAccs
      [VName]
    infusible
      (Let Pat (LetDec SOACS)
    pat1 StmAux (ExpDec SOACS)
    aux1 (WithAcc [WithAccInput SOACS]
    w_inps1 Lambda SOACS
    lam1))
      (Let Pat (LetDec SOACS)
    pat2 StmAux (ExpDec SOACS)
    aux2 (WithAcc [WithAccInput SOACS]
    w_inps2 Lambda SOACS
    lam2))
        | ([PatElem Type]
    pat1_els, [PatElem Type]
    pat2_els) <- (Pat Type -> [PatElem Type]
    forall dec. Pat dec -> [PatElem dec]
    patElems Pat Type
    Pat (LetDec SOACS)
    pat1, Pat Type -> [PatElem Type]
    forall dec. Pat dec -> [PatElem dec]
    patElems Pat Type
    Pat (LetDec SOACS)
    pat2),
          ([AccTup]
    acc_tup1, [(PatElem (LetDec SOACS), SubExpRes)]
    other_pr1) <- [PatElem (LetDec SOACS)]
    -> [WithAccInput SOACS]
    -> Lambda SOACS
    -> ([AccTup], [(PatElem (LetDec SOACS), SubExpRes)])
    groupAccs [PatElem Type]
    [PatElem (LetDec SOACS)]
    pat1_els [WithAccInput SOACS]
    w_inps1 Lambda SOACS
    lam1,
          ([AccTup]
    acc_tup2, [(PatElem (LetDec SOACS), SubExpRes)]
    other_pr2) <- [PatElem (LetDec SOACS)]
    -> [WithAccInput SOACS]
    -> Lambda SOACS
    -> ([AccTup], [(PatElem (LetDec SOACS), SubExpRes)])
    groupAccs [PatElem Type]
    [PatElem (LetDec SOACS)]
    pat2_els [WithAccInput SOACS]
    w_inps2 Lambda SOACS
    lam2,
          ([(AccTup, AccTup)]
    tup_common, [AccTup]
    acc_tup1', [AccTup]
    acc_tup2') <-
            [AccTup] -> [AccTup] -> ([(AccTup, AccTup)], [AccTup], [AccTup])
    groupCommonAccs [AccTup]
    acc_tup1 [AccTup]
    acc_tup2,
          -- safety 0: make sure that the accs from acc_tup1' and
          --           acc_tup2' do not overlap
          [VName]
    pnms_1' <- (PatElem Type -> VName) -> [PatElem Type] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map PatElem Type -> VName
    forall dec. PatElem dec -> VName
    patElemName ([PatElem Type] -> [VName]) -> [PatElem Type] -> [VName]
    forall a b. (a -> b) -> a -> b
    $ (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))
     -> [PatElem Type])
    -> [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))]
    -> [PatElem Type]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap (\([PatElem Type]
    nms, WithAccInput SOACS
    _, Param Type
    _, Param Type
    _, (VName, Certs)
    _) -> [PatElem Type]
    nms) [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    [AccTup]
    acc_tup1',
          [VName]
    winp_2' <- (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))
     -> [VName])
    -> [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))]
    -> [VName]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap (\([PatElem Type]
    _, (ShapeBase SubExp
    _, [VName]
    nms, Maybe (Lambda SOACS, [SubExp])
    _), Param Type
    _, Param Type
    _, (VName, Certs)
    _) -> [VName]
    nms) [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    [AccTup]
    acc_tup2',
          Bool -> Bool
    not (Bool -> Bool) -> Bool -> Bool
    forall a b. (a -> b) -> a -> b
    $ Names -> Names -> Bool
    namesIntersect ([VName] -> Names
    namesFromList [VName]
    pnms_1') ([VName] -> Names
    namesFromList [VName]
    winp_2'),
          -- safety 1: we have already determined the commons;
          --           now we also need to check NOT-IN FV(lam2)
          Bool -> Bool
    not (Bool -> Bool) -> Bool -> Bool
    forall a b. (a -> b) -> a -> b
    $ Names -> Names -> Bool
    namesIntersect ([VName] -> Names
    namesFromList [VName]
    pnms_1') (Lambda SOACS -> Names
    forall a. FreeIn a => a -> Names
    freeIn Lambda SOACS
    lam2),
          -- safety 2:
          -- bs <- map patElemName $ concatMap accTup1 acc_tup1,
          [VName]
    bs <- (PatElem Type -> VName) -> [PatElem Type] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map PatElem Type -> VName
    forall dec. PatElem dec -> VName
    patElemName ([PatElem Type] -> [VName]) -> [PatElem Type] -> [VName]
    forall a b. (a -> b) -> a -> b
    $ ((([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)),
      ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)))
     -> [PatElem Type])
    -> [(([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)),
         ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)))]
    -> [PatElem Type]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
     (VName, Certs))
    -> [PatElem Type]
    AccTup -> [PatElem (LetDec SOACS)]
    accTup1 (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))
     -> [PatElem Type])
    -> ((([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)),
         ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)))
        -> ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
            (VName, Certs)))
    -> (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs)),
        ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs)))
    -> [PatElem Type]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs)),
     ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs)))
    -> ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
        (VName, Certs))
    forall a b. (a, b) -> a
    fst) [(([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)),
      ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)))]
    [(AccTup, AccTup)]
    tup_common,
          (VName -> Bool) -> [VName] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all (VName -> [VName] -> Bool
    forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
    `notElem` [VName]
    infusible) [VName]
    bs,
          -- safety 3:
          Names
    cs <- [VName] -> Names
    namesFromList ([VName] -> Names) -> [VName] -> Names
    forall a b. (a -> b) -> a -> b
    $ (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))
     -> [VName])
    -> [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))]
    -> [VName]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap ((\(ShapeBase SubExp
    _, [VName]
    xs, Maybe (Lambda SOACS, [SubExp])
    _) -> [VName]
    xs) (WithAccInput SOACS -> [VName])
    -> (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))
        -> WithAccInput SOACS)
    -> ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
        (VName, Certs))
    -> [VName]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
     (VName, Certs))
    -> WithAccInput SOACS
    AccTup -> WithAccInput SOACS
    accTup2) [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    [AccTup]
    acc_tup2,
          ((PatElem Type, SubExpRes) -> Bool)
    -> [(PatElem Type, SubExpRes)] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all ((VName -> Names -> Bool
    `notNameIn` Names
    cs) (VName -> Bool)
    -> ((PatElem Type, SubExpRes) -> VName)
    -> (PatElem Type, SubExpRes)
    -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . PatElem Type -> VName
    forall dec. PatElem dec -> VName
    patElemName (PatElem Type -> VName)
    -> ((PatElem Type, SubExpRes) -> PatElem Type)
    -> (PatElem Type, SubExpRes)
    -> VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (PatElem Type, SubExpRes) -> PatElem Type
    forall a b. (a, b) -> a
    fst) [(PatElem Type, SubExpRes)]
    [(PatElem (LetDec SOACS), SubExpRes)]
    other_pr1 = m (Stm SOACS) -> Maybe (m (Stm SOACS))
    forall a. a -> Maybe a
    Just (m (Stm SOACS) -> Maybe (m (Stm SOACS)))
    -> m (Stm SOACS) -> Maybe (m (Stm SOACS))
    forall a b. (a -> b) -> a -> b
    $ do
            let getCertPairs :: (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs)),
     ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs)))
    -> (VName, VName)
    getCertPairs (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
     (VName, Certs))
    t1, ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
     (VName, Certs))
    t2) = (Param Type -> VName
    forall dec. Param dec -> VName
    paramName (AccTup -> LParam SOACS
    accTup3 ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
     (VName, Certs))
    AccTup
    t2), Param Type -> VName
    forall dec. Param dec -> VName
    paramName (AccTup -> LParam SOACS
    accTup3 ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
     (VName, Certs))
    AccTup
    t1))
                tab_certs :: Map VName VName
    tab_certs = [(VName, VName)] -> Map VName VName
    forall k a. Ord k => [(k, a)] -> Map k a
    M.fromList ([(VName, VName)] -> Map VName VName)
    -> [(VName, VName)] -> Map VName VName
    forall a b. (a -> b) -> a -> b
    $ ((([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)),
      ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)))
     -> (VName, VName))
    -> [(([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)),
         ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)))]
    -> [(VName, VName)]
    forall a b. (a -> b) -> [a] -> [b]
    map (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs)),
     ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs)))
    -> (VName, VName)
    getCertPairs [(([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)),
      ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)))]
    [(AccTup, AccTup)]
    tup_common
                lam2_bdy' :: Body SOACS
    lam2_bdy' = Map VName VName -> Body SOACS -> Body SOACS
    forall a. Substitute a => Map VName VName -> a -> a
    substituteNames Map VName VName
    tab_certs (Lambda SOACS -> Body SOACS
    forall rep. Lambda rep -> Body rep
    lambdaBody Lambda SOACS
    lam2)
                rcrt_params :: [Param Type]
    rcrt_params = ((([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)),
      ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)))
     -> Param Type)
    -> [(([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)),
         ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)))]
    -> [Param Type]
    forall a b. (a -> b) -> [a] -> [b]
    map (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
     (VName, Certs))
    -> Param Type
    AccTup -> LParam SOACS
    accTup3 (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))
     -> Param Type)
    -> ((([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)),
         ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)))
        -> ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
            (VName, Certs)))
    -> (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs)),
        ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs)))
    -> Param Type
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs)),
     ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs)))
    -> ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
        (VName, Certs))
    forall a b. (a, b) -> a
    fst) [(([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)),
      ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)))]
    [(AccTup, AccTup)]
    tup_common [Param Type] -> [Param Type] -> [Param Type]
    forall a. [a] -> [a] -> [a]
    ++ (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))
     -> Param Type)
    -> [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))]
    -> [Param Type]
    forall a b. (a -> b) -> [a] -> [b]
    map ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
     (VName, Certs))
    -> Param Type
    AccTup -> LParam SOACS
    accTup3 [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    [AccTup]
    acc_tup1' [Param Type] -> [Param Type] -> [Param Type]
    forall a. [a] -> [a] -> [a]
    ++ (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))
     -> Param Type)
    -> [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))]
    -> [Param Type]
    forall a b. (a -> b) -> [a] -> [b]
    map ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
     (VName, Certs))
    -> Param Type
    AccTup -> LParam SOACS
    accTup3 [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    [AccTup]
    acc_tup2'
                racc_params :: [Param Type]
    racc_params = ((([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)),
      ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)))
     -> Param Type)
    -> [(([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)),
         ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)))]
    -> [Param Type]
    forall a b. (a -> b) -> [a] -> [b]
    map (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
     (VName, Certs))
    -> Param Type
    AccTup -> LParam SOACS
    accTup4 (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))
     -> Param Type)
    -> ((([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)),
         ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)))
        -> ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
            (VName, Certs)))
    -> (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs)),
        ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs)))
    -> Param Type
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs)),
     ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs)))
    -> ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
        (VName, Certs))
    forall a b. (a, b) -> a
    fst) [(([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)),
      ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)))]
    [(AccTup, AccTup)]
    tup_common [Param Type] -> [Param Type] -> [Param Type]
    forall a. [a] -> [a] -> [a]
    ++ (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))
     -> Param Type)
    -> [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))]
    -> [Param Type]
    forall a b. (a -> b) -> [a] -> [b]
    map ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
     (VName, Certs))
    -> Param Type
    AccTup -> LParam SOACS
    accTup4 [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    [AccTup]
    acc_tup1' [Param Type] -> [Param Type] -> [Param Type]
    forall a. [a] -> [a] -> [a]
    ++ (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))
     -> Param Type)
    -> [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))]
    -> [Param Type]
    forall a b. (a -> b) -> [a] -> [b]
    map ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
     (VName, Certs))
    -> Param Type
    AccTup -> LParam SOACS
    accTup4 [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    [AccTup]
    acc_tup2'
                ([VName]
    comm_res_nms, [Certs]
    comm_res_certs2) = [(VName, Certs)] -> ([VName], [Certs])
    forall a b. [(a, b)] -> ([a], [b])
    unzip ([(VName, Certs)] -> ([VName], [Certs]))
    -> [(VName, Certs)] -> ([VName], [Certs])
    forall a b. (a -> b) -> a -> b
    $ ((([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)),
      ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)))
     -> (VName, Certs))
    -> [(([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)),
         ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)))]
    -> [(VName, Certs)]
    forall a b. (a -> b) -> [a] -> [b]
    map (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
     (VName, Certs))
    -> (VName, Certs)
    AccTup -> (VName, Certs)
    accTup5 (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))
     -> (VName, Certs))
    -> ((([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)),
         ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)))
        -> ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
            (VName, Certs)))
    -> (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs)),
        ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs)))
    -> (VName, Certs)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs)),
     ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs)))
    -> ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
        (VName, Certs))
    forall a b. (a, b) -> b
    snd) [(([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)),
      ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)))]
    [(AccTup, AccTup)]
    tup_common
                ([VName]
    _, [Certs]
    comm_res_certs1) = [(VName, Certs)] -> ([VName], [Certs])
    forall a b. [(a, b)] -> ([a], [b])
    unzip ([(VName, Certs)] -> ([VName], [Certs]))
    -> [(VName, Certs)] -> ([VName], [Certs])
    forall a b. (a -> b) -> a -> b
    $ ((([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)),
      ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)))
     -> (VName, Certs))
    -> [(([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)),
         ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)))]
    -> [(VName, Certs)]
    forall a b. (a -> b) -> [a] -> [b]
    map (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
     (VName, Certs))
    -> (VName, Certs)
    AccTup -> (VName, Certs)
    accTup5 (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))
     -> (VName, Certs))
    -> ((([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)),
         ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)))
        -> ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
            (VName, Certs)))
    -> (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs)),
        ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs)))
    -> (VName, Certs)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs)),
     ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs)))
    -> ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
        (VName, Certs))
    forall a b. (a, b) -> a
    fst) [(([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)),
      ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)))]
    [(AccTup, AccTup)]
    tup_common
                com_res_certs :: [Certs]
    com_res_certs = (Certs -> Certs -> Certs) -> [Certs] -> [Certs] -> [Certs]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith (\Certs
    x Certs
    y -> [VName] -> Certs
    Certs (Certs -> [VName]
    unCerts Certs
    x [VName] -> [VName] -> [VName]
    forall a. [a] -> [a] -> [a]
    ++ Certs -> [VName]
    unCerts Certs
    y)) [Certs]
    comm_res_certs1 [Certs]
    comm_res_certs2
                bdyres_certs :: [Certs]
    bdyres_certs = [Certs]
    com_res_certs [Certs] -> [Certs] -> [Certs]
    forall a. [a] -> [a] -> [a]
    ++ (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))
     -> Certs)
    -> [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))]
    -> [Certs]
    forall a b. (a -> b) -> [a] -> [b]
    map ((VName, Certs) -> Certs
    forall a b. (a, b) -> b
    snd ((VName, Certs) -> Certs)
    -> (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))
        -> (VName, Certs))
    -> ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
        (VName, Certs))
    -> Certs
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
     (VName, Certs))
    -> (VName, Certs)
    AccTup -> (VName, Certs)
    accTup5) ([([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    [AccTup]
    acc_tup1' [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    -> [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))]
    -> [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))]
    forall a. [a] -> [a] -> [a]
    ++ [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    [AccTup]
    acc_tup2')
                bdyres_accse :: [SubExp]
    bdyres_accse = (VName -> SubExp) -> [VName] -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map VName -> SubExp
    Var [VName]
    comm_res_nms [SubExp] -> [SubExp] -> [SubExp]
    forall a. [a] -> [a] -> [a]
    ++ (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))
     -> SubExp)
    -> [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))]
    -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map (VName -> SubExp
    Var (VName -> SubExp)
    -> (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))
        -> VName)
    -> ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
        (VName, Certs))
    -> SubExp
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (VName, Certs) -> VName
    forall a b. (a, b) -> a
    fst ((VName, Certs) -> VName)
    -> (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))
        -> (VName, Certs))
    -> ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
        (VName, Certs))
    -> VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
     (VName, Certs))
    -> (VName, Certs)
    AccTup -> (VName, Certs)
    accTup5) ([([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    [AccTup]
    acc_tup1' [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    -> [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))]
    -> [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))]
    forall a. [a] -> [a] -> [a]
    ++ [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    [AccTup]
    acc_tup2')
                bdy_res_accs :: Result
    bdy_res_accs = (Certs -> SubExp -> SubExpRes) -> [Certs] -> [SubExp] -> Result
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith Certs -> SubExp -> SubExpRes
    SubExpRes [Certs]
    bdyres_certs [SubExp]
    bdyres_accse
                bdy_res_others :: Result
    bdy_res_others = ((PatElem Type, SubExpRes) -> SubExpRes)
    -> [(PatElem Type, SubExpRes)] -> Result
    forall a b. (a -> b) -> [a] -> [b]
    map (PatElem Type, SubExpRes) -> SubExpRes
    forall a b. (a, b) -> b
    snd ([(PatElem Type, SubExpRes)] -> Result)
    -> [(PatElem Type, SubExpRes)] -> Result
    forall a b. (a -> b) -> a -> b
    $ [(PatElem Type, SubExpRes)]
    [(PatElem (LetDec SOACS), SubExpRes)]
    other_pr1 [(PatElem Type, SubExpRes)]
    -> [(PatElem Type, SubExpRes)] -> [(PatElem Type, SubExpRes)]
    forall a. [a] -> [a] -> [a]
    ++ [(PatElem Type, SubExpRes)]
    [(PatElem (LetDec SOACS), SubExpRes)]
    other_pr2
            scope <- m (Scope SOACS)
    forall rep (m :: * -> *). HasScope rep m => m (Scope rep)
    askScope
            lam_bdy <-
              runBodyBuilder $ do
                localScope (scope <> scopeOfLParams (rcrt_params ++ racc_params)) $ do
                  -- add the stms of lam1
                  mapM_ addStm $ stmsToList $ bodyStms $ lambdaBody lam1
                  -- add the copy stms for the common accumulator
                  forM_ tup_common $ \(([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
     (VName, Certs))
    tup1, ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
     (VName, Certs))
    tup2) -> do
                    let (Param Type
    lpar1, Param Type
    lpar2) = (AccTup -> LParam SOACS
    accTup4 ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
     (VName, Certs))
    AccTup
    tup1, AccTup -> LParam SOACS
    accTup4 ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
     (VName, Certs))
    AccTup
    tup2)
                        ((VName
    nm1, Certs
    _), VName
    nm2, Type
    tp_acc) = (AccTup -> (VName, Certs)
    accTup5 ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
     (VName, Certs))
    AccTup
    tup1, Param Type -> VName
    forall dec. Param dec -> VName
    paramName Param Type
    lpar2, Param Type -> Type
    forall dec. Param dec -> dec
    paramDec Param Type
    lpar1)
                    Pat (LetDec (Rep (BuilderT SOACS (State VNameSource))))
    -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    -> BuilderT SOACS (State VNameSource) ()
    forall (m :: * -> *).
    MonadBuilder m =>
    Pat (LetDec (Rep m)) -> Exp (Rep m) -> m ()
    letBind ([PatElem Type] -> Pat Type
    forall dec. [PatElem dec] -> Pat dec
    Pat [VName -> Type -> PatElem Type
    forall dec. VName -> dec -> PatElem dec
    PatElem VName
    nm2 Type
    tp_acc]) (Exp (Rep (BuilderT SOACS (State VNameSource)))
     -> BuilderT SOACS (State VNameSource) ())
    -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    -> BuilderT SOACS (State VNameSource) ()
    forall a b. (a -> b) -> a -> b
    $ BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    forall rep. BasicOp -> Exp rep
    BasicOp (BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource))))
    -> BasicOp -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    forall a b. (a -> b) -> a -> b
    $ SubExp -> BasicOp
    SubExp (SubExp -> BasicOp) -> SubExp -> BasicOp
    forall a b. (a -> b) -> a -> b
    $ VName -> SubExp
    Var VName
    nm1
                  -- add copy stms to bring in scope x1 ... xq
                  forM_ other_pr1 $ \(PatElem Type
    pat_elm, SubExpRes
    bdy_res) -> do
                    let (VName
    nm, SubExp
    se, Type
    tp) = (PatElem Type -> VName
    forall dec. PatElem dec -> VName
    patElemName PatElem Type
    pat_elm, SubExpRes -> SubExp
    resSubExp SubExpRes
    bdy_res, PatElem Type -> Type
    forall dec. Typed dec => PatElem dec -> Type
    patElemType PatElem Type
    pat_elm)
                    Certs
    -> BuilderT SOACS (State VNameSource) ()
    -> BuilderT SOACS (State VNameSource) ()
    forall a.
    Certs
    -> BuilderT SOACS (State VNameSource) a
    -> BuilderT SOACS (State VNameSource) a
    forall (m :: * -> *) a. MonadBuilder m => Certs -> m a -> m a
    certifying (SubExpRes -> Certs
    resCerts SubExpRes
    bdy_res) (BuilderT SOACS (State VNameSource) ()
     -> BuilderT SOACS (State VNameSource) ())
    -> BuilderT SOACS (State VNameSource) ()
    -> BuilderT SOACS (State VNameSource) ()
    forall a b. (a -> b) -> a -> b
    $
                      Pat (LetDec (Rep (BuilderT SOACS (State VNameSource))))
    -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    -> BuilderT SOACS (State VNameSource) ()
    forall (m :: * -> *).
    MonadBuilder m =>
    Pat (LetDec (Rep m)) -> Exp (Rep m) -> m ()
    letBind ([PatElem Type] -> Pat Type
    forall dec. [PatElem dec] -> Pat dec
    Pat [VName -> Type -> PatElem Type
    forall dec. VName -> dec -> PatElem dec
    PatElem VName
    nm Type
    tp]) (Exp (Rep (BuilderT SOACS (State VNameSource)))
     -> BuilderT SOACS (State VNameSource) ())
    -> Exp (Rep (BuilderT SOACS (State VNameSource)))
    -> BuilderT SOACS (State VNameSource) ()
    forall a b. (a -> b) -> a -> b
    $
                        BasicOp -> Exp SOACS
    forall rep. BasicOp -> Exp rep
    BasicOp (SubExp -> BasicOp
    SubExp SubExp
    se)
                  -- add the statements of lam2 (in which the acc-certificates have been substituted)
                  mapM_ addStm $ stmsToList $ bodyStms lam2_bdy'
                  -- build the result of body
                  pure $ bdy_res_accs ++ bdy_res_others
            let tp_res_other = ((PatElem Type, SubExpRes) -> Type)
    -> [(PatElem Type, SubExpRes)] -> [Type]
    forall a b. (a -> b) -> [a] -> [b]
    map (PatElem Type -> Type
    forall dec. Typed dec => PatElem dec -> Type
    patElemType (PatElem Type -> Type)
    -> ((PatElem Type, SubExpRes) -> PatElem Type)
    -> (PatElem Type, SubExpRes)
    -> Type
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (PatElem Type, SubExpRes) -> PatElem Type
    forall a b. (a, b) -> a
    fst) ([(PatElem Type, SubExpRes)]
    [(PatElem (LetDec SOACS), SubExpRes)]
    other_pr1 [(PatElem Type, SubExpRes)]
    -> [(PatElem Type, SubExpRes)] -> [(PatElem Type, SubExpRes)]
    forall a. [a] -> [a] -> [a]
    ++ [(PatElem Type, SubExpRes)]
    [(PatElem (LetDec SOACS), SubExpRes)]
    other_pr2)
                res_lam =
                  Lambda
                    { lambdaParams :: [LParam SOACS]
    lambdaParams = [Param Type]
    rcrt_params [Param Type] -> [Param Type] -> [Param Type]
    forall a. [a] -> [a] -> [a]
    ++ [Param Type]
    racc_params,
                      lambdaBody :: Body SOACS
    lambdaBody = Body SOACS
    lam_bdy,
                      lambdaReturnType :: [Type]
    lambdaReturnType = (Param Type -> Type) -> [Param Type] -> [Type]
    forall a b. (a -> b) -> [a] -> [b]
    map Param Type -> Type
    forall dec. Param dec -> dec
    paramDec [Param Type]
    racc_params [Type] -> [Type] -> [Type]
    forall a. [a] -> [a] -> [a]
    ++ [Type]
    tp_res_other
                    }
            res_lam' <- renameLambda res_lam
            let res_pat =
                  ((([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)),
      ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)))
     -> [PatElem Type])
    -> [(([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)),
         ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)))]
    -> [PatElem Type]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
     (VName, Certs))
    -> [PatElem Type]
    AccTup -> [PatElem (LetDec SOACS)]
    accTup1 (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))
     -> [PatElem Type])
    -> ((([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)),
         ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)))
        -> ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
            (VName, Certs)))
    -> (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs)),
        ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs)))
    -> [PatElem Type]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs)),
     ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs)))
    -> ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
        (VName, Certs))
    forall a b. (a, b) -> b
    snd) [(([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)),
      ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)))]
    [(AccTup, AccTup)]
    tup_common
                    [PatElem Type] -> [PatElem Type] -> [PatElem Type]
    forall a. [a] -> [a] -> [a]
    ++ (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))
     -> [PatElem Type])
    -> [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))]
    -> [PatElem Type]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
     (VName, Certs))
    -> [PatElem Type]
    AccTup -> [PatElem (LetDec SOACS)]
    accTup1 ([([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    [AccTup]
    acc_tup1' [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    -> [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))]
    -> [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))]
    forall a. [a] -> [a] -> [a]
    ++ [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    [AccTup]
    acc_tup2')
                    [PatElem Type] -> [PatElem Type] -> [PatElem Type]
    forall a. [a] -> [a] -> [a]
    ++ ((PatElem Type, SubExpRes) -> PatElem Type)
    -> [(PatElem Type, SubExpRes)] -> [PatElem Type]
    forall a b. (a -> b) -> [a] -> [b]
    map (PatElem Type, SubExpRes) -> PatElem Type
    forall a b. (a, b) -> a
    fst ([(PatElem Type, SubExpRes)]
    [(PatElem (LetDec SOACS), SubExpRes)]
    other_pr1 [(PatElem Type, SubExpRes)]
    -> [(PatElem Type, SubExpRes)] -> [(PatElem Type, SubExpRes)]
    forall a. [a] -> [a] -> [a]
    ++ [(PatElem Type, SubExpRes)]
    [(PatElem (LetDec SOACS), SubExpRes)]
    other_pr2)
                res_w_inps = ((([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)),
      ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)))
     -> WithAccInput SOACS)
    -> [(([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)),
         ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)))]
    -> [WithAccInput SOACS]
    forall a b. (a -> b) -> [a] -> [b]
    map (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
     (VName, Certs))
    -> WithAccInput SOACS
    AccTup -> WithAccInput SOACS
    accTup2 (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))
     -> WithAccInput SOACS)
    -> ((([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)),
         ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)))
        -> ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
            (VName, Certs)))
    -> (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs)),
        ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs)))
    -> WithAccInput SOACS
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs)),
     ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs)))
    -> ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
        (VName, Certs))
    forall a b. (a, b) -> a
    fst) [(([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)),
      ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)))]
    [(AccTup, AccTup)]
    tup_common [WithAccInput SOACS]
    -> [WithAccInput SOACS] -> [WithAccInput SOACS]
    forall a. [a] -> [a] -> [a]
    ++ (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))
     -> WithAccInput SOACS)
    -> [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))]
    -> [WithAccInput SOACS]
    forall a b. (a -> b) -> [a] -> [b]
    map ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
     (VName, Certs))
    -> WithAccInput SOACS
    AccTup -> WithAccInput SOACS
    accTup2 ([([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    [AccTup]
    acc_tup1' [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    -> [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))]
    -> [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))]
    forall a. [a] -> [a] -> [a]
    ++ [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    [AccTup]
    acc_tup2')
            res_w_inps' <- mapM renameLamInWAccInp res_w_inps
            pure $ Let (Pat res_pat) (aux1 <> aux2) $ WithAcc res_w_inps' res_lam'
        where
          -- local helpers:

          groupAccs ::
            [PatElem (LetDec SOACS)] ->
            [WithAccInput SOACS] ->
            Lambda SOACS ->
            ([AccTup], [(PatElem (LetDec SOACS), SubExpRes)])
          groupAccs :: [PatElem (LetDec SOACS)]
    -> [WithAccInput SOACS]
    -> Lambda SOACS
    -> ([AccTup], [(PatElem (LetDec SOACS), SubExpRes)])
    groupAccs [PatElem (LetDec SOACS)]
    pat_els [WithAccInput SOACS]
    wacc_inps Lambda SOACS
    wlam =
            let lam_params :: [LParam SOACS]
    lam_params = Lambda SOACS -> [LParam SOACS]
    forall rep. Lambda rep -> [LParam rep]
    lambdaParams Lambda SOACS
    wlam
                n :: Int
    n = [Param Type] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [Param Type]
    lam_params
                ([Param Type]
    lam_par_crts, [Param Type]
    lam_par_accs) = Int -> [Param Type] -> ([Param Type], [Param Type])
    forall a. Int -> [a] -> ([a], [a])
    splitAt (Int
    n Int -> Int -> Int
    forall a. Integral a => a -> a -> a
    `div` Int
    2) [Param Type]
    lam_params
                lab_res_ses :: Result
    lab_res_ses = Body SOACS -> Result
    forall rep. Body rep -> Result
    bodyResult (Body SOACS -> Result) -> Body SOACS -> Result
    forall a b. (a -> b) -> a -> b
    $ Lambda SOACS -> Body SOACS
    forall rep. Lambda rep -> Body rep
    lambdaBody Lambda SOACS
    wlam
             in [PatElem (LetDec SOACS)]
    -> [WithAccInput SOACS]
    -> [LParam SOACS]
    -> [LParam SOACS]
    -> Result
    -> ([AccTup], [(PatElem (LetDec SOACS), SubExpRes)])
    groupAccsHlp [PatElem (LetDec SOACS)]
    pat_els [WithAccInput SOACS]
    wacc_inps [Param Type]
    [LParam SOACS]
    lam_par_crts [Param Type]
    [LParam SOACS]
    lam_par_accs Result
    lab_res_ses
          groupAccsHlp ::
            [PatElem (LetDec SOACS)] ->
            [WithAccInput SOACS] ->
            [LParam SOACS] ->
            [LParam SOACS] ->
            [SubExpRes] ->
            ([AccTup], [(PatElem (LetDec SOACS), SubExpRes)])
          groupAccsHlp :: [PatElem (LetDec SOACS)]
    -> [WithAccInput SOACS]
    -> [LParam SOACS]
    -> [LParam SOACS]
    -> Result
    -> ([AccTup], [(PatElem (LetDec SOACS), SubExpRes)])
    groupAccsHlp [PatElem (LetDec SOACS)]
    pat_els [] [] [] Result
    lam_res_ses
            | [PatElem Type] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [PatElem Type]
    [PatElem (LetDec SOACS)]
    pat_els Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Result -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length Result
    lam_res_ses =
                ([], [PatElem Type] -> Result -> [(PatElem Type, SubExpRes)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [PatElem Type]
    [PatElem (LetDec SOACS)]
    pat_els Result
    lam_res_ses)
          groupAccsHlp
            [PatElem (LetDec SOACS)]
    pat_els
            (winp :: WithAccInput SOACS
    winp@(ShapeBase SubExp
    _, [VName]
    inp, Maybe (Lambda SOACS, [SubExp])
    _) : [WithAccInput SOACS]
    wacc_inps)
            (LParam SOACS
    par_crt : [LParam SOACS]
    lam_par_crts)
            (LParam SOACS
    par_acc : [LParam SOACS]
    lam_par_accs)
            (SubExpRes
    res_se : Result
    lam_res_ses)
              | Int
    n <- [VName] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [VName]
    inp,
                (Int
    n Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    <= [PatElem Type] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [PatElem Type]
    [PatElem (LetDec SOACS)]
    pat_els) Bool -> Bool -> Bool
    && (Int
    n Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    <= (Int
    1 Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Result -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length Result
    lam_res_ses)),
                Var VName
    res_nm <- SubExpRes -> SubExp
    resSubExp SubExpRes
    res_se =
                  let ([PatElem Type]
    pat_els_cur, [PatElem Type]
    pat_els') = Int -> [PatElem Type] -> ([PatElem Type], [PatElem Type])
    forall a. Int -> [a] -> ([a], [a])
    splitAt Int
    n [PatElem Type]
    [PatElem (LetDec SOACS)]
    pat_els
                      ([AccTup]
    rec1, [(PatElem (LetDec SOACS), SubExpRes)]
    rec2) = [PatElem (LetDec SOACS)]
    -> [WithAccInput SOACS]
    -> [LParam SOACS]
    -> [LParam SOACS]
    -> Result
    -> ([AccTup], [(PatElem (LetDec SOACS), SubExpRes)])
    groupAccsHlp [PatElem Type]
    [PatElem (LetDec SOACS)]
    pat_els' [WithAccInput SOACS]
    wacc_inps [LParam SOACS]
    lam_par_crts [LParam SOACS]
    lam_par_accs Result
    lam_res_ses
                   in (([PatElem Type]
    pat_els_cur, WithAccInput SOACS
    winp, Param Type
    LParam SOACS
    par_crt, Param Type
    LParam SOACS
    par_acc, (VName
    res_nm, SubExpRes -> Certs
    resCerts SubExpRes
    res_se)) ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
     (VName, Certs))
    -> [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))]
    -> [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))]
    forall a. a -> [a] -> [a]
    : [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    rec1, [(PatElem Type, SubExpRes)]
    [(PatElem (LetDec SOACS), SubExpRes)]
    rec2)
          groupAccsHlp [PatElem (LetDec SOACS)]
    _ [WithAccInput SOACS]
    _ [LParam SOACS]
    _ [LParam SOACS]
    _ Result
    _ =
            String
    -> ([([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs))],
        [(PatElem Type, SubExpRes)])
    forall a. HasCallStack => String -> a
    error String
    "Unreachable case reached in groupAccsHlp!"
          --
          groupCommonAccs :: [AccTup] -> [AccTup] -> ([(AccTup, AccTup)], [AccTup], [AccTup])
          groupCommonAccs :: [AccTup] -> [AccTup] -> ([(AccTup, AccTup)], [AccTup], [AccTup])
    groupCommonAccs [] [AccTup]
    tup_accs2 =
            ([], [], [AccTup]
    tup_accs2)
          groupCommonAccs (AccTup
    tup_acc1 : [AccTup]
    tup_accs1) [AccTup]
    tup_accs2
            | [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    commons2 <- (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))
     -> Bool)
    -> [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))]
    -> [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))]
    forall a. (a -> Bool) -> [a] -> [a]
    filter (AccTup -> AccTup -> Bool
    matchingAccTup AccTup
    tup_acc1) [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    [AccTup]
    tup_accs2,
              [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    commons2 Int -> Int -> Bool
    forall a. Ord a => a -> a -> Bool
    <= Int
    1 =
                let ([(AccTup, AccTup)]
    rec1, [AccTup]
    rec2, [AccTup]
    rec3) =
                      [AccTup] -> [AccTup] -> ([(AccTup, AccTup)], [AccTup], [AccTup])
    groupCommonAccs [AccTup]
    tup_accs1 ([AccTup] -> ([(AccTup, AccTup)], [AccTup], [AccTup]))
    -> [AccTup] -> ([(AccTup, AccTup)], [AccTup], [AccTup])
    forall a b. (a -> b) -> a -> b
    $
                        if [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    -> Bool
    forall a. [a] -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    commons2
                          then [AccTup]
    tup_accs2
                          else (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))
     -> Bool)
    -> [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))]
    -> [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))]
    forall a. (a -> Bool) -> [a] -> [a]
    filter (Bool -> Bool
    not (Bool -> Bool)
    -> (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))
        -> Bool)
    -> ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
        (VName, Certs))
    -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . AccTup -> AccTup -> Bool
    matchingAccTup AccTup
    tup_acc1) [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    [AccTup]
    tup_accs2
                 in if [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    -> Bool
    forall a. [a] -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    commons2
                      then ([(([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)),
      ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)))]
    [(AccTup, AccTup)]
    rec1, ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
     (VName, Certs))
    AccTup
    tup_acc1 ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
     (VName, Certs))
    -> [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))]
    -> [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
         (VName, Certs))]
    forall a. a -> [a] -> [a]
    : [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    rec2, [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    [AccTup]
    rec3)
                      else ((([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
     (VName, Certs))
    AccTup
    tup_acc1, [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    -> ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
        (VName, Certs))
    forall a. HasCallStack => [a] -> a
    head [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    commons2) (([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs)),
     ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs)))
    -> [(([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)),
         ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)))]
    -> [(([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)),
         ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs)))]
    forall a. a -> [a] -> [a]
    : [(([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)),
      ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
       (VName, Certs)))]
    rec1, [AccTup]
    tup_accs1, [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
      (VName, Certs))]
    [AccTup]
    rec3)
          groupCommonAccs [AccTup]
    _ [AccTup]
    _ =
            String
    -> ([(([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
           (VName, Certs)),
          ([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
           (VName, Certs)))],
        [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs))],
        [([PatElem Type], WithAccInput SOACS, Param Type, Param Type,
          (VName, Certs))])
    forall a. HasCallStack => String -> a
    error String
    "Unreachable case reached in groupCommonAccs!"
          renameLamInWAccInp :: (a, b, Maybe (Lambda rep, b)) -> m (a, b, Maybe (Lambda rep, b))
    renameLamInWAccInp (a
    shp, b
    inps, Just (Lambda rep
    lam, b
    se)) = do
            lam' <- Lambda rep -> m (Lambda rep)
    forall rep (m :: * -> *).
    (Renameable rep, MonadFreshNames m) =>
    Lambda rep -> m (Lambda rep)
    renameLambda Lambda rep
    lam
            pure (shp, inps, Just (lam', se))
          renameLamInWAccInp (a, b, Maybe (Lambda rep, b))
    winp = (a, b, Maybe (Lambda rep, b)) -> m (a, b, Maybe (Lambda rep, b))
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (a, b, Maybe (Lambda rep, b))
    winp
    --
    tryFuseWithAccs [VName]
    _ Stm SOACS
    _ Stm SOACS
    _ =
      Maybe (m (Stm SOACS))
    forall a. Maybe a
    Nothing

    -------------------------------
    --- simple helper functions ---
    -------------------------------

    equivLambda ::
      M.Map VName VName ->
      Lambda SOACS ->
      Lambda SOACS ->
      Bool
    equivLambda :: Map VName VName -> Lambda SOACS -> Lambda SOACS -> Bool
    equivLambda Map VName VName
    stab Lambda SOACS
    lam1 Lambda SOACS
    lam2
      | ([Param Type]
    ps1, [Param Type]
    ps2) <- (Lambda SOACS -> [LParam SOACS]
    forall rep. Lambda rep -> [LParam rep]
    lambdaParams Lambda SOACS
    lam1, Lambda SOACS -> [LParam SOACS]
    forall rep. Lambda rep -> [LParam rep]
    lambdaParams Lambda SOACS
    lam2),
        ([VName]
    nms1, [VName]
    nms2) <- ((Param Type -> VName) -> [Param Type] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map Param Type -> VName
    forall dec. Param dec -> VName
    paramName [Param Type]
    ps1, (Param Type -> VName) -> [Param Type] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map Param Type -> VName
    forall dec. Param dec -> VName
    paramName [Param Type]
    ps2),
        (Param Type -> Type) -> [Param Type] -> [Type]
    forall a b. (a -> b) -> [a] -> [b]
    map Param Type -> Type
    forall dec. Param dec -> dec
    paramDec [Param Type]
    ps1 [Type] -> [Type] -> Bool
    forall a. Eq a => a -> a -> Bool
    == (Param Type -> Type) -> [Param Type] -> [Type]
    forall a b. (a -> b) -> [a] -> [b]
    map Param Type -> Type
    forall dec. Param dec -> dec
    paramDec [Param Type]
    ps2,
        (Param Type -> Attrs) -> [Param Type] -> [Attrs]
    forall a b. (a -> b) -> [a] -> [b]
    map Param Type -> Attrs
    forall dec. Param dec -> Attrs
    paramAttrs [Param Type]
    ps1 [Attrs] -> [Attrs] -> Bool
    forall a. Eq a => a -> a -> Bool
    == (Param Type -> Attrs) -> [Param Type] -> [Attrs]
    forall a b. (a -> b) -> [a] -> [b]
    map Param Type -> Attrs
    forall dec. Param dec -> Attrs
    paramAttrs [Param Type]
    ps2,
        Lambda SOACS -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda SOACS
    lam1 [Type] -> [Type] -> Bool
    forall a. Eq a => a -> a -> Bool
    == Lambda SOACS -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda SOACS
    lam2,
        (Body SOACS
    bdy1, Body SOACS
    bdy2) <- (Lambda SOACS -> Body SOACS
    forall rep. Lambda rep -> Body rep
    lambdaBody Lambda SOACS
    lam1, Lambda SOACS -> Body SOACS
    forall rep. Lambda rep -> Body rep
    lambdaBody Lambda SOACS
    lam2),
        Body SOACS -> BodyDec SOACS
    forall rep. Body rep -> BodyDec rep
    bodyDec Body SOACS
    bdy1 () -> () -> Bool
    forall a. Eq a => a -> a -> Bool
    == Body SOACS -> BodyDec SOACS
    forall rep. Body rep -> BodyDec rep
    bodyDec Body SOACS
    bdy2 =
          let insert :: Map k a -> (a, k) -> Map k a
    insert Map k a
    tab (a
    x, k
    k) = k -> a -> Map k a -> Map k a
    forall k a. Ord k => k -> a -> Map k a -> Map k a
    M.insert k
    k a
    x Map k a
    tab
              stab' :: Map VName VName
    stab' = (Map VName VName -> (VName, VName) -> Map VName VName)
    -> Map VName VName -> [(VName, VName)] -> Map VName VName
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl Map VName VName -> (VName, VName) -> Map VName VName
    forall {k} {a}. Ord k => Map k a -> (a, k) -> Map k a
    insert Map VName VName
    stab ([(VName, VName)] -> Map VName VName)
    -> [(VName, VName)] -> Map VName VName
    forall a b. (a -> b) -> a -> b
    $ [VName] -> [VName] -> [(VName, VName)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [VName]
    nms1 [VName]
    nms2
              fStm :: (Map VName VName, Bool)
    -> (Stm SOACS, Stm SOACS) -> (Map VName VName, Bool)
    fStm (Map VName VName
    vtab, Bool
    False) (Stm SOACS, Stm SOACS)
    _ = (Map VName VName
    vtab, Bool
    False)
              fStm (Map VName VName
    vtab, Bool
    True) (Stm SOACS
    s1, Stm SOACS
    s2) = Map VName VName
    -> Stm SOACS -> Stm SOACS -> (Map VName VName, Bool)
    equivStm Map VName VName
    vtab Stm SOACS
    s1 Stm SOACS
    s2
              (Map VName VName
    stab'', Bool
    success) =
                ((Map VName VName, Bool)
     -> (Stm SOACS, Stm SOACS) -> (Map VName VName, Bool))
    -> (Map VName VName, Bool)
    -> [(Stm SOACS, Stm SOACS)]
    -> (Map VName VName, Bool)
    forall b a. (b -> a -> b) -> b -> [a] -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl (Map VName VName, Bool)
    -> (Stm SOACS, Stm SOACS) -> (Map VName VName, Bool)
    fStm (Map VName VName
    stab', Bool
    True) ([(Stm SOACS, Stm SOACS)] -> (Map VName VName, Bool))
    -> [(Stm SOACS, Stm SOACS)] -> (Map VName VName, Bool)
    forall a b. (a -> b) -> a -> b
    $
                  [Stm SOACS] -> [Stm SOACS] -> [(Stm SOACS, Stm SOACS)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip (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
    bdy1)) ([Stm SOACS] -> [(Stm SOACS, Stm SOACS)])
    -> [Stm SOACS] -> [(Stm SOACS, Stm SOACS)]
    forall a b. (a -> b) -> a -> b
    $
                    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
    bdy2)
              sres2 :: [SubExp]
    sres2 = Map VName VName -> [SubExp] -> [SubExp]
    substInSEs Map VName VName
    stab'' ([SubExp] -> [SubExp]) -> [SubExp] -> [SubExp]
    forall a b. (a -> b) -> a -> b
    $ (SubExpRes -> SubExp) -> Result -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExpRes -> SubExp
    resSubExp (Result -> [SubExp]) -> Result -> [SubExp]
    forall a b. (a -> b) -> a -> b
    $ Body SOACS -> Result
    forall rep. Body rep -> Result
    bodyResult Body SOACS
    bdy2
           in Bool
    success Bool -> Bool -> Bool
    && (SubExpRes -> SubExp) -> Result -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExpRes -> SubExp
    resSubExp (Body SOACS -> Result
    forall rep. Body rep -> Result
    bodyResult Body SOACS
    bdy1) [SubExp] -> [SubExp] -> Bool
    forall a. Eq a => a -> a -> Bool
    == [SubExp]
    sres2
    equivLambda Map VName VName
    _ Lambda SOACS
    _ Lambda SOACS
    _ =
      Bool
    False

    equivStm ::
      M.Map VName VName ->
      Stm SOACS ->
      Stm SOACS ->
      (M.Map VName VName, Bool)
    equivStm :: Map VName VName
    -> Stm SOACS -> Stm SOACS -> (Map VName VName, Bool)
    equivStm
      Map VName VName
    stab
      (Let Pat (LetDec SOACS)
    pat1 StmAux (ExpDec SOACS)
    aux1 (BasicOp (BinOp BinOp
    bop1 SubExp
    se11 SubExp
    se12)))
      (Let Pat (LetDec SOACS)
    pat2 StmAux (ExpDec SOACS)
    aux2 (BasicOp (BinOp BinOp
    bop2 SubExp
    se21 SubExp
    se22)))
        | [SubExp
    se11, SubExp
    se12] [SubExp] -> [SubExp] -> Bool
    forall a. Eq a => a -> a -> Bool
    == Map VName VName -> [SubExp] -> [SubExp]
    substInSEs Map VName VName
    stab [SubExp
    se21, SubExp
    se22],
          ([PatElem Type]
    pels1, [PatElem Type]
    pels2) <- (Pat Type -> [PatElem Type]
    forall dec. Pat dec -> [PatElem dec]
    patElems Pat Type
    Pat (LetDec SOACS)
    pat1, Pat Type -> [PatElem Type]
    forall dec. Pat dec -> [PatElem dec]
    patElems Pat Type
    Pat (LetDec SOACS)
    pat2),
          (PatElem Type -> Type) -> [PatElem Type] -> [Type]
    forall a b. (a -> b) -> [a] -> [b]
    map PatElem Type -> Type
    forall dec. PatElem dec -> dec
    patElemDec [PatElem Type]
    pels1 [Type] -> [Type] -> Bool
    forall a. Eq a => a -> a -> Bool
    == (PatElem Type -> Type) -> [PatElem Type] -> [Type]
    forall a b. (a -> b) -> [a] -> [b]
    map PatElem Type -> Type
    forall dec. PatElem dec -> dec
    patElemDec [PatElem Type]
    pels2,
          BinOp
    bop1 BinOp -> BinOp -> Bool
    forall a. Eq a => a -> a -> Bool
    == BinOp
    bop2 Bool -> Bool -> Bool
    && StmAux ()
    StmAux (ExpDec SOACS)
    aux1 StmAux () -> StmAux () -> Bool
    forall a. Eq a => a -> a -> Bool
    == StmAux ()
    StmAux (ExpDec SOACS)
    aux2 =
            let stab_new :: Map VName VName
    stab_new =
                  [(VName, VName)] -> Map VName VName
    forall k a. Ord k => [(k, a)] -> Map k a
    M.fromList ([(VName, VName)] -> Map VName VName)
    -> [(VName, VName)] -> Map VName VName
    forall a b. (a -> b) -> a -> b
    $
                    [VName] -> [VName] -> [(VName, VName)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip ((PatElem Type -> VName) -> [PatElem Type] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map PatElem Type -> VName
    forall dec. PatElem dec -> VName
    patElemName [PatElem Type]
    pels2) ((PatElem Type -> VName) -> [PatElem Type] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map PatElem Type -> VName
    forall dec. PatElem dec -> VName
    patElemName [PatElem Type]
    pels1)
             in (Map VName VName -> Map VName VName -> Map VName VName
    forall k a. Ord k => Map k a -> Map k a -> Map k a
    M.union Map VName VName
    stab_new Map VName VName
    stab, Bool
    True)
    -- To Be Continued ...
    equivStm Map VName VName
    vtab Stm SOACS
    _ Stm SOACS
    _ = (Map VName VName
    vtab, Bool
    False)

    matchingAccTup :: AccTup -> AccTup -> Bool
    matchingAccTup :: AccTup -> AccTup -> Bool
    matchingAccTup
      ([PatElem (LetDec SOACS)]
    pat_els1, (ShapeBase SubExp
    shp1, [VName]
    _winp_arrs1, Maybe (Lambda SOACS, [SubExp])
    mlam1), LParam SOACS
    _, LParam SOACS
    _, (VName, Certs)
    _)
      ([PatElem (LetDec SOACS)]
    _, (ShapeBase SubExp
    shp2, [VName]
    winp_arrs2, Maybe (Lambda SOACS, [SubExp])
    mlam2), LParam SOACS
    _, LParam SOACS
    _, (VName, Certs)
    _) =
        ShapeBase SubExp -> [SubExp]
    forall d. ShapeBase d -> [d]
    shapeDims ShapeBase SubExp
    shp1 [SubExp] -> [SubExp] -> Bool
    forall a. Eq a => a -> a -> Bool
    == ShapeBase SubExp -> [SubExp]
    forall d. ShapeBase d -> [d]
    shapeDims ShapeBase SubExp
    shp2
          Bool -> Bool -> Bool
    && (PatElem Type -> VName) -> [PatElem Type] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map PatElem Type -> VName
    forall dec. PatElem dec -> VName
    patElemName [PatElem Type]
    [PatElem (LetDec SOACS)]
    pat_els1 [VName] -> [VName] -> Bool
    forall a. Eq a => a -> a -> Bool
    == [VName]
    winp_arrs2
          Bool -> Bool -> Bool
    && case (Maybe (Lambda SOACS, [SubExp])
    mlam1, Maybe (Lambda SOACS, [SubExp])
    mlam2) of
            (Maybe (Lambda SOACS, [SubExp])
    Nothing, Maybe (Lambda SOACS, [SubExp])
    Nothing) -> Bool
    True
            (Just (Lambda SOACS
    lam1, [SubExp]
    see1), Just (Lambda SOACS
    lam2, [SubExp]
    see2)) ->
              ([SubExp]
    see1 [SubExp] -> [SubExp] -> Bool
    forall a. Eq a => a -> a -> Bool
    == [SubExp]
    see2) Bool -> Bool -> Bool
    && Map VName VName -> Lambda SOACS -> Lambda SOACS -> Bool
    equivLambda Map VName VName
    forall k a. Map k a
    M.empty Lambda SOACS
    lam1 Lambda SOACS
    lam2
            (Maybe (Lambda SOACS, [SubExp]), Maybe (Lambda SOACS, [SubExp]))
    _ -> Bool
    False

    substInSEs :: M.Map VName VName -> [SubExp] -> [SubExp]
    substInSEs :: Map VName VName -> [SubExp] -> [SubExp]
    substInSEs Map VName VName
    vtab = (SubExp -> SubExp) -> [SubExp] -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExp -> SubExp
    substInSE
      where
        substInSE :: SubExp -> SubExp
    substInSE (Var VName
    x)
          | Just VName
    y <- VName -> Map VName VName -> Maybe VName
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    x Map VName VName
    vtab = VName -> SubExp
    Var VName
    y
        substInSE SubExp
    z = SubExp
    z