    {-# LANGUAGE TypeFamilies #-}
    {-# LANGUAGE UndecidableInstances #-}

    -- | Definition of /Second-Order Array Combinators/ (SOACs), which are
    -- the main form of parallelism in the early stages of the compiler.
    module Futhark.IR.SOACS.SOAC
      ( SOAC (..),
        ScremaForm (..),
        ScatterSpec,
        HistOp (..),
        Scan (..),
        scanResults,
        singleScan,
        Reduce (..),
        redResults,
        singleReduce,

        -- * Utility
        scremaType,
        soacType,
        typeCheckSOAC,
        mkIdentityLambda,
        isIdentityLambda,
        nilFn,
        scanomapSOAC,
        redomapSOAC,
        scanSOAC,
        reduceSOAC,
        mapSOAC,
        isScanomapSOAC,
        isRedomapSOAC,
        isScanSOAC,
        isReduceSOAC,
        isMapSOAC,
        ppScrema,
        ppHist,
        ppStream,
        ppScatter,
        groupScatterResults,
        groupScatterResults',
        splitScatterResults,

        -- * Generic traversal
        SOACMapper (..),
        identitySOACMapper,
        mapSOACM,
        traverseSOACStms,
      )
    where

    import Control.Category
    import Control.Monad
    import Control.Monad.Identity
    import Control.Monad.State.Strict
    import Control.Monad.Writer
    import Data.Function ((&))
    import Data.List (intersperse)
    import Data.Map.Strict qualified as M
    import Data.Maybe
    import Futhark.Analysis.Alias qualified as Alias
    import Futhark.Analysis.DataDependencies
    import Futhark.Analysis.Metrics
    import Futhark.Analysis.PrimExp.Convert
    import Futhark.Analysis.SymbolTable qualified as ST
    import Futhark.Construct
    import Futhark.IR
    import Futhark.IR.Aliases (Aliases, CanBeAliased (..))
    import Futhark.IR.Prop.Aliases
    import Futhark.IR.TypeCheck qualified as TC
    import Futhark.Optimise.Simplify.Rep
    import Futhark.Transform.Rename
    import Futhark.Transform.Substitute
    import Futhark.Util (chunks, maybeNth, splitAt3)
    import Futhark.Util.Pretty (Doc, align, comma, commasep, docText, parens, ppTuple', pretty, (<+>), (</>))
    import Futhark.Util.Pretty qualified as PP
    import Prelude hiding (id, (.))

    -- | How the results of a scatter operation should be written. Each
    -- element of the list consists of a @v@ (often a `VName`) specifying
    -- which array to scatter to, a `Shape` describing the shape of that
    -- array, and an `Int` describing how many elements should be written
    -- to that array for each invocation of the scatter.
    type ScatterSpec v = [(Shape, Int, v)]

    -- | A second-order array combinator (SOAC).
    data SOAC rep
      = Stream SubExp [VName] [SubExp] (Lambda rep)
      | -- | @Scatter <length> <inputs> <lambda> <spec>@
        --
        -- Scatter maps values from a set of input arrays to indices and values of a
        -- set of output arrays. It is able to write multiple values to multiple
        -- outputs each of which may have multiple dimensions.
        --
        -- <inputs> is a list of input arrays, all having size <length>, elements of
        -- which are applied to the <lambda> function. For instance, if there are
        -- two arrays, <lambda> will get two values as input, one from each array.
        --
        -- <spec> specifies the result of the <lambda> and which arrays to
        -- write to.
        --
        -- <lambda> is a function that takes inputs from <inputs> and
        -- returns values according to <spec>. It returns values in the
        -- following manner:
        --
        --     [index_0, index_1, ..., index_n, value_0, value_1, ..., value_m]
        --
        -- For each output in <spec>, <lambda> returns <i> * <j> index
        -- values and <j> output values, where <i> is the number of
        -- dimensions (rank) of the given output, and <j> is the number of
        -- output values written to the given output.
        --
        -- For example, given the following scatter specification:
        --
        --     [([x1, y1, z1], 2, arr1), ([x2, y2], 1, arr2)]
        --
        -- <lambda> will produce 6 (3 * 2) index values and 2 output values for
        -- <arr1>, and 2 (2 * 1) index values and 1 output value for
        -- arr2. Additionally, the results are grouped, so the first 6 index values
        -- will correspond to the first two output values, and so on. For this
        -- example, <lambda> should return a total of 11 values, 8 index values and
        -- 3 output values.  See also 'splitScatterResults'.
        Scatter SubExp [VName] (ScatterSpec VName) (Lambda rep)
      | -- | @Hist <length> <input arrays> <dest-arrays-and-ops> <bucket fun>@
        --
        -- The final lambda produces indexes and values for the 'HistOp's.
        Hist SubExp [VName] [HistOp rep] (Lambda rep)
      | -- FIXME: this should not be here
        JVP [SubExp] [SubExp] (Lambda rep)
      | -- FIXME: this should not be here
        VJP [SubExp] [SubExp] (Lambda rep)
      | -- | A combination of scan, reduction, and map.  The first
        -- t'SubExp' is the size of the input arrays.
        Screma SubExp [VName] (ScremaForm rep)
      deriving (SOAC rep -> SOAC rep -> Bool
    (SOAC rep -> SOAC rep -> Bool)
    -> (SOAC rep -> SOAC rep -> Bool) -> Eq (SOAC rep)
    forall rep. RepTypes rep => SOAC rep -> SOAC rep -> Bool
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: forall rep. RepTypes rep => SOAC rep -> SOAC rep -> Bool
    == :: SOAC rep -> SOAC rep -> Bool
    $c/= :: forall rep. RepTypes rep => SOAC rep -> SOAC rep -> Bool
    /= :: SOAC rep -> SOAC rep -> Bool
    Eq, Eq (SOAC rep)
    Eq (SOAC rep) =>
    (SOAC rep -> SOAC rep -> Ordering)
    -> (SOAC rep -> SOAC rep -> Bool)
    -> (SOAC rep -> SOAC rep -> Bool)
    -> (SOAC rep -> SOAC rep -> Bool)
    -> (SOAC rep -> SOAC rep -> Bool)
    -> (SOAC rep -> SOAC rep -> SOAC rep)
    -> (SOAC rep -> SOAC rep -> SOAC rep)
    -> Ord (SOAC rep)
    SOAC rep -> SOAC rep -> Bool
    SOAC rep -> SOAC rep -> Ordering
    SOAC rep -> SOAC rep -> SOAC rep
    forall a.
    Eq a =>
    (a -> a -> Ordering)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> a)
    -> (a -> a -> a)
    -> Ord a
    forall rep. RepTypes rep => Eq (SOAC rep)
    forall rep. RepTypes rep => SOAC rep -> SOAC rep -> Bool
    forall rep. RepTypes rep => SOAC rep -> SOAC rep -> Ordering
    forall rep. RepTypes rep => SOAC rep -> SOAC rep -> SOAC rep
    $ccompare :: forall rep. RepTypes rep => SOAC rep -> SOAC rep -> Ordering
    compare :: SOAC rep -> SOAC rep -> Ordering
    $c< :: forall rep. RepTypes rep => SOAC rep -> SOAC rep -> Bool
    < :: SOAC rep -> SOAC rep -> Bool
    $c<= :: forall rep. RepTypes rep => SOAC rep -> SOAC rep -> Bool
    <= :: SOAC rep -> SOAC rep -> Bool
    $c> :: forall rep. RepTypes rep => SOAC rep -> SOAC rep -> Bool
    > :: SOAC rep -> SOAC rep -> Bool
    $c>= :: forall rep. RepTypes rep => SOAC rep -> SOAC rep -> Bool
    >= :: SOAC rep -> SOAC rep -> Bool
    $cmax :: forall rep. RepTypes rep => SOAC rep -> SOAC rep -> SOAC rep
    max :: SOAC rep -> SOAC rep -> SOAC rep
    $cmin :: forall rep. RepTypes rep => SOAC rep -> SOAC rep -> SOAC rep
    min :: SOAC rep -> SOAC rep -> SOAC rep
    Ord, Int -> SOAC rep -> ShowS
    [SOAC rep] -> ShowS
    SOAC rep -> String
    (Int -> SOAC rep -> ShowS)
    -> (SOAC rep -> String) -> ([SOAC rep] -> ShowS) -> Show (SOAC rep)
    forall rep. RepTypes rep => Int -> SOAC rep -> ShowS
    forall rep. RepTypes rep => [SOAC rep] -> ShowS
    forall rep. RepTypes rep => SOAC rep -> String
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: forall rep. RepTypes rep => Int -> SOAC rep -> ShowS
    showsPrec :: Int -> SOAC rep -> ShowS
    $cshow :: forall rep. RepTypes rep => SOAC rep -> String
    show :: SOAC rep -> String
    $cshowList :: forall rep. RepTypes rep => [SOAC rep] -> ShowS
    showList :: [SOAC rep] -> ShowS
    Show)

    -- | Information about computing a single histogram.
    data HistOp rep = HistOp
      { forall rep. HistOp rep -> Shape
    histShape :: Shape,
        -- | Race factor @RF@ means that only @1/RF@
        -- bins are used.
        forall rep. HistOp rep -> SubExp
    histRaceFactor :: SubExp,
        forall rep. HistOp rep -> [VName]
    histDest :: [VName],
        forall rep. HistOp rep -> [SubExp]
    histNeutral :: [SubExp],
        forall rep. HistOp rep -> Lambda rep
    histOp :: Lambda rep
      }
      deriving (HistOp rep -> HistOp rep -> Bool
    (HistOp rep -> HistOp rep -> Bool)
    -> (HistOp rep -> HistOp rep -> Bool) -> Eq (HistOp rep)
    forall rep. RepTypes rep => HistOp rep -> HistOp rep -> Bool
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: forall rep. RepTypes rep => HistOp rep -> HistOp rep -> Bool
    == :: HistOp rep -> HistOp rep -> Bool
    $c/= :: forall rep. RepTypes rep => HistOp rep -> HistOp rep -> Bool
    /= :: HistOp rep -> HistOp rep -> Bool
    Eq, Eq (HistOp rep)
    Eq (HistOp rep) =>
    (HistOp rep -> HistOp rep -> Ordering)
    -> (HistOp rep -> HistOp rep -> Bool)
    -> (HistOp rep -> HistOp rep -> Bool)
    -> (HistOp rep -> HistOp rep -> Bool)
    -> (HistOp rep -> HistOp rep -> Bool)
    -> (HistOp rep -> HistOp rep -> HistOp rep)
    -> (HistOp rep -> HistOp rep -> HistOp rep)
    -> Ord (HistOp rep)
    HistOp rep -> HistOp rep -> Bool
    HistOp rep -> HistOp rep -> Ordering
    HistOp rep -> HistOp rep -> HistOp rep
    forall a.
    Eq a =>
    (a -> a -> Ordering)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> a)
    -> (a -> a -> a)
    -> Ord a
    forall rep. RepTypes rep => Eq (HistOp rep)
    forall rep. RepTypes rep => HistOp rep -> HistOp rep -> Bool
    forall rep. RepTypes rep => HistOp rep -> HistOp rep -> Ordering
    forall rep. RepTypes rep => HistOp rep -> HistOp rep -> HistOp rep
    $ccompare :: forall rep. RepTypes rep => HistOp rep -> HistOp rep -> Ordering
    compare :: HistOp rep -> HistOp rep -> Ordering
    $c< :: forall rep. RepTypes rep => HistOp rep -> HistOp rep -> Bool
    < :: HistOp rep -> HistOp rep -> Bool
    $c<= :: forall rep. RepTypes rep => HistOp rep -> HistOp rep -> Bool
    <= :: HistOp rep -> HistOp rep -> Bool
    $c> :: forall rep. RepTypes rep => HistOp rep -> HistOp rep -> Bool
    > :: HistOp rep -> HistOp rep -> Bool
    $c>= :: forall rep. RepTypes rep => HistOp rep -> HistOp rep -> Bool
    >= :: HistOp rep -> HistOp rep -> Bool
    $cmax :: forall rep. RepTypes rep => HistOp rep -> HistOp rep -> HistOp rep
    max :: HistOp rep -> HistOp rep -> HistOp rep
    $cmin :: forall rep. RepTypes rep => HistOp rep -> HistOp rep -> HistOp rep
    min :: HistOp rep -> HistOp rep -> HistOp rep
    Ord, Int -> HistOp rep -> ShowS
    [HistOp rep] -> ShowS
    HistOp rep -> String
    (Int -> HistOp rep -> ShowS)
    -> (HistOp rep -> String)
    -> ([HistOp rep] -> ShowS)
    -> Show (HistOp rep)
    forall rep. RepTypes rep => Int -> HistOp rep -> ShowS
    forall rep. RepTypes rep => [HistOp rep] -> ShowS
    forall rep. RepTypes rep => HistOp rep -> String
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: forall rep. RepTypes rep => Int -> HistOp rep -> ShowS
    showsPrec :: Int -> HistOp rep -> ShowS
    $cshow :: forall rep. RepTypes rep => HistOp rep -> String
    show :: HistOp rep -> String
    $cshowList :: forall rep. RepTypes rep => [HistOp rep] -> ShowS
    showList :: [HistOp rep] -> ShowS
    Show)

    -- | The essential parts of a 'Screma' factored out (everything
    -- except the input arrays).
    data ScremaForm rep = ScremaForm
      { -- | The "main" lambda of the Screma. For a map, this is
        -- equivalent to 'isMapSOAC'. Note that the meaning of the return
        -- value of this lambda depends crucially on exactly which Screma
        -- this is. The parameters will correspond exactly to elements of
        -- the input arrays, however.
        forall rep. ScremaForm rep -> Lambda rep
    scremaLambda :: Lambda rep,
        forall rep. ScremaForm rep -> [Scan rep]
    scremaScans :: [Scan rep],
        forall rep. ScremaForm rep -> [Reduce rep]
    scremaReduces :: [Reduce rep]
      }
      deriving (ScremaForm rep -> ScremaForm rep -> Bool
    (ScremaForm rep -> ScremaForm rep -> Bool)
    -> (ScremaForm rep -> ScremaForm rep -> Bool)
    -> Eq (ScremaForm rep)
    forall rep.
    RepTypes rep =>
    ScremaForm rep -> ScremaForm rep -> Bool
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: forall rep.
    RepTypes rep =>
    ScremaForm rep -> ScremaForm rep -> Bool
    == :: ScremaForm rep -> ScremaForm rep -> Bool
    $c/= :: forall rep.
    RepTypes rep =>
    ScremaForm rep -> ScremaForm rep -> Bool
    /= :: ScremaForm rep -> ScremaForm rep -> Bool
    Eq, Eq (ScremaForm rep)
    Eq (ScremaForm rep) =>
    (ScremaForm rep -> ScremaForm rep -> Ordering)
    -> (ScremaForm rep -> ScremaForm rep -> Bool)
    -> (ScremaForm rep -> ScremaForm rep -> Bool)
    -> (ScremaForm rep -> ScremaForm rep -> Bool)
    -> (ScremaForm rep -> ScremaForm rep -> Bool)
    -> (ScremaForm rep -> ScremaForm rep -> ScremaForm rep)
    -> (ScremaForm rep -> ScremaForm rep -> ScremaForm rep)
    -> Ord (ScremaForm rep)
    ScremaForm rep -> ScremaForm rep -> Bool
    ScremaForm rep -> ScremaForm rep -> Ordering
    ScremaForm rep -> ScremaForm rep -> ScremaForm rep
    forall a.
    Eq a =>
    (a -> a -> Ordering)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> a)
    -> (a -> a -> a)
    -> Ord a
    forall rep. RepTypes rep => Eq (ScremaForm rep)
    forall rep.
    RepTypes rep =>
    ScremaForm rep -> ScremaForm rep -> Bool
    forall rep.
    RepTypes rep =>
    ScremaForm rep -> ScremaForm rep -> Ordering
    forall rep.
    RepTypes rep =>
    ScremaForm rep -> ScremaForm rep -> ScremaForm rep
    $ccompare :: forall rep.
    RepTypes rep =>
    ScremaForm rep -> ScremaForm rep -> Ordering
    compare :: ScremaForm rep -> ScremaForm rep -> Ordering
    $c< :: forall rep.
    RepTypes rep =>
    ScremaForm rep -> ScremaForm rep -> Bool
    < :: ScremaForm rep -> ScremaForm rep -> Bool
    $c<= :: forall rep.
    RepTypes rep =>
    ScremaForm rep -> ScremaForm rep -> Bool
    <= :: ScremaForm rep -> ScremaForm rep -> Bool
    $c> :: forall rep.
    RepTypes rep =>
    ScremaForm rep -> ScremaForm rep -> Bool
    > :: ScremaForm rep -> ScremaForm rep -> Bool
    $c>= :: forall rep.
    RepTypes rep =>
    ScremaForm rep -> ScremaForm rep -> Bool
    >= :: ScremaForm rep -> ScremaForm rep -> Bool
    $cmax :: forall rep.
    RepTypes rep =>
    ScremaForm rep -> ScremaForm rep -> ScremaForm rep
    max :: ScremaForm rep -> ScremaForm rep -> ScremaForm rep
    $cmin :: forall rep.
    RepTypes rep =>
    ScremaForm rep -> ScremaForm rep -> ScremaForm rep
    min :: ScremaForm rep -> ScremaForm rep -> ScremaForm rep
    Ord, Int -> ScremaForm rep -> ShowS
    [ScremaForm rep] -> ShowS
    ScremaForm rep -> String
    (Int -> ScremaForm rep -> ShowS)
    -> (ScremaForm rep -> String)
    -> ([ScremaForm rep] -> ShowS)
    -> Show (ScremaForm rep)
    forall rep. RepTypes rep => Int -> ScremaForm rep -> ShowS
    forall rep. RepTypes rep => [ScremaForm rep] -> ShowS
    forall rep. RepTypes rep => ScremaForm rep -> String
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: forall rep. RepTypes rep => Int -> ScremaForm rep -> ShowS
    showsPrec :: Int -> ScremaForm rep -> ShowS
    $cshow :: forall rep. RepTypes rep => ScremaForm rep -> String
    show :: ScremaForm rep -> String
    $cshowList :: forall rep. RepTypes rep => [ScremaForm rep] -> ShowS
    showList :: [ScremaForm rep] -> ShowS
    Show)

    singleBinOp :: (Buildable rep) => [Lambda rep] -> Lambda rep
    singleBinOp :: forall rep. Buildable rep => [Lambda rep] -> Lambda rep
    singleBinOp [Lambda rep]
    lams =
      Lambda
        { lambdaParams :: [LParam rep]
    lambdaParams = (Lambda rep -> [Param Type]) -> [Lambda rep] -> [Param Type]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap Lambda rep -> [Param Type]
    Lambda rep -> [LParam rep]
    forall {rep}. Lambda rep -> [Param (LParamInfo rep)]
    xParams [Lambda rep]
    lams [Param Type] -> [Param Type] -> [Param Type]
    forall a. [a] -> [a] -> [a]
    ++ (Lambda rep -> [Param Type]) -> [Lambda rep] -> [Param Type]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap Lambda rep -> [Param Type]
    Lambda rep -> [LParam rep]
    forall {rep}. Lambda rep -> [Param (LParamInfo rep)]
    yParams [Lambda rep]
    lams,
          lambdaReturnType :: [Type]
    lambdaReturnType = (Lambda rep -> [Type]) -> [Lambda rep] -> [Type]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap Lambda rep -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType [Lambda rep]
    lams,
          lambdaBody :: Body rep
    lambdaBody =
            Stms rep -> Result -> Body rep
    forall rep. Buildable rep => Stms rep -> Result -> Body rep
    mkBody
              ([Stms rep] -> Stms rep
    forall a. Monoid a => [a] -> a
    mconcat ((Lambda rep -> Stms rep) -> [Lambda rep] -> [Stms rep]
    forall a b. (a -> b) -> [a] -> [b]
    map (Body rep -> Stms rep
    forall rep. Body rep -> Stms rep
    bodyStms (Body rep -> Stms rep)
    -> (Lambda rep -> Body rep) -> Lambda rep -> Stms rep
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Lambda rep -> Body rep
    forall rep. Lambda rep -> Body rep
    lambdaBody) [Lambda rep]
    lams))
              ((Lambda rep -> Result) -> [Lambda rep] -> Result
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap (Body rep -> Result
    forall rep. Body rep -> Result
    bodyResult (Body rep -> Result)
    -> (Lambda rep -> Body rep) -> Lambda rep -> Result
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Lambda rep -> Body rep
    forall rep. Lambda rep -> Body rep
    lambdaBody) [Lambda rep]
    lams)
        }
      where
        xParams :: Lambda rep -> [Param (LParamInfo rep)]
    xParams Lambda rep
    lam = Int -> [Param (LParamInfo rep)] -> [Param (LParamInfo rep)]
    forall a. Int -> [a] -> [a]
    take ([Type] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length (Lambda rep -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda rep
    lam)) (Lambda rep -> [Param (LParamInfo rep)]
    forall {rep}. Lambda rep -> [Param (LParamInfo rep)]
    lambdaParams Lambda rep
    lam)
        yParams :: Lambda rep -> [Param (LParamInfo rep)]
    yParams Lambda rep
    lam = Int -> [Param (LParamInfo rep)] -> [Param (LParamInfo rep)]
    forall a. Int -> [a] -> [a]
    drop ([Type] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length (Lambda rep -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda rep
    lam)) (Lambda rep -> [Param (LParamInfo rep)]
    forall {rep}. Lambda rep -> [Param (LParamInfo rep)]
    lambdaParams Lambda rep
    lam)

    -- | How to compute a single scan result.
    data Scan rep = Scan
      { forall rep. Scan rep -> Lambda rep
    scanLambda :: Lambda rep,
        forall rep. Scan rep -> [SubExp]
    scanNeutral :: [SubExp]
      }
      deriving (Scan rep -> Scan rep -> Bool
    (Scan rep -> Scan rep -> Bool)
    -> (Scan rep -> Scan rep -> Bool) -> Eq (Scan rep)
    forall rep. RepTypes rep => Scan rep -> Scan rep -> Bool
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: forall rep. RepTypes rep => Scan rep -> Scan rep -> Bool
    == :: Scan rep -> Scan rep -> Bool
    $c/= :: forall rep. RepTypes rep => Scan rep -> Scan rep -> Bool
    /= :: Scan rep -> Scan rep -> Bool
    Eq, Eq (Scan rep)
    Eq (Scan rep) =>
    (Scan rep -> Scan rep -> Ordering)
    -> (Scan rep -> Scan rep -> Bool)
    -> (Scan rep -> Scan rep -> Bool)
    -> (Scan rep -> Scan rep -> Bool)
    -> (Scan rep -> Scan rep -> Bool)
    -> (Scan rep -> Scan rep -> Scan rep)
    -> (Scan rep -> Scan rep -> Scan rep)
    -> Ord (Scan rep)
    Scan rep -> Scan rep -> Bool
    Scan rep -> Scan rep -> Ordering
    Scan rep -> Scan rep -> Scan rep
    forall a.
    Eq a =>
    (a -> a -> Ordering)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> a)
    -> (a -> a -> a)
    -> Ord a
    forall rep. RepTypes rep => Eq (Scan rep)
    forall rep. RepTypes rep => Scan rep -> Scan rep -> Bool
    forall rep. RepTypes rep => Scan rep -> Scan rep -> Ordering
    forall rep. RepTypes rep => Scan rep -> Scan rep -> Scan rep
    $ccompare :: forall rep. RepTypes rep => Scan rep -> Scan rep -> Ordering
    compare :: Scan rep -> Scan rep -> Ordering
    $c< :: forall rep. RepTypes rep => Scan rep -> Scan rep -> Bool
    < :: Scan rep -> Scan rep -> Bool
    $c<= :: forall rep. RepTypes rep => Scan rep -> Scan rep -> Bool
    <= :: Scan rep -> Scan rep -> Bool
    $c> :: forall rep. RepTypes rep => Scan rep -> Scan rep -> Bool
    > :: Scan rep -> Scan rep -> Bool
    $c>= :: forall rep. RepTypes rep => Scan rep -> Scan rep -> Bool
    >= :: Scan rep -> Scan rep -> Bool
    $cmax :: forall rep. RepTypes rep => Scan rep -> Scan rep -> Scan rep
    max :: Scan rep -> Scan rep -> Scan rep
    $cmin :: forall rep. RepTypes rep => Scan rep -> Scan rep -> Scan rep
    min :: Scan rep -> Scan rep -> Scan rep
    Ord, Int -> Scan rep -> ShowS
    [Scan rep] -> ShowS
    Scan rep -> String
    (Int -> Scan rep -> ShowS)
    -> (Scan rep -> String) -> ([Scan rep] -> ShowS) -> Show (Scan rep)
    forall rep. RepTypes rep => Int -> Scan rep -> ShowS
    forall rep. RepTypes rep => [Scan rep] -> ShowS
    forall rep. RepTypes rep => Scan rep -> String
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: forall rep. RepTypes rep => Int -> Scan rep -> ShowS
    showsPrec :: Int -> Scan rep -> ShowS
    $cshow :: forall rep. RepTypes rep => Scan rep -> String
    show :: Scan rep -> String
    $cshowList :: forall rep. RepTypes rep => [Scan rep] -> ShowS
    showList :: [Scan rep] -> ShowS
    Show)

    -- | What are the sizes of reduction results produced by these 'Scan's?
    scanSizes :: [Scan rep] -> [Int]
    scanSizes :: forall rep. [Scan rep] -> [Int]
    scanSizes = (Scan rep -> Int) -> [Scan rep] -> [Int]
    forall a b. (a -> b) -> [a] -> [b]
    map ([SubExp] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length ([SubExp] -> Int) -> (Scan rep -> [SubExp]) -> Scan rep -> Int
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Scan rep -> [SubExp]
    forall rep. Scan rep -> [SubExp]
    scanNeutral)

    -- | How many reduction results are produced by these 'Scan's?
    scanResults :: [Scan rep] -> Int
    scanResults :: forall rep. [Scan rep] -> Int
    scanResults = [Int] -> Int
    forall a. Num a => [a] -> a
    forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
    sum ([Int] -> Int) -> ([Scan rep] -> [Int]) -> [Scan rep] -> Int
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . [Scan rep] -> [Int]
    forall rep. [Scan rep] -> [Int]
    scanSizes

    -- | Combine multiple scan operators to a single operator.
    singleScan :: (Buildable rep) => [Scan rep] -> Scan rep
    singleScan :: forall rep. Buildable rep => [Scan rep] -> Scan rep
    singleScan [Scan rep]
    scans =
      let scan_nes :: [SubExp]
    scan_nes = (Scan rep -> [SubExp]) -> [Scan rep] -> [SubExp]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap Scan rep -> [SubExp]
    forall rep. Scan rep -> [SubExp]
    scanNeutral [Scan rep]
    scans
          scan_lam :: Lambda rep
    scan_lam = [Lambda rep] -> Lambda rep
    forall rep. Buildable rep => [Lambda rep] -> Lambda rep
    singleBinOp ([Lambda rep] -> Lambda rep) -> [Lambda rep] -> Lambda rep
    forall a b. (a -> b) -> a -> b
    $ (Scan rep -> Lambda rep) -> [Scan rep] -> [Lambda rep]
    forall a b. (a -> b) -> [a] -> [b]
    map Scan rep -> Lambda rep
    forall rep. Scan rep -> Lambda rep
    scanLambda [Scan rep]
    scans
       in Lambda rep -> [SubExp] -> Scan rep
    forall rep. Lambda rep -> [SubExp] -> Scan rep
    Scan Lambda rep
    scan_lam [SubExp]
    scan_nes

    -- | How to compute a single reduction result.
    data Reduce rep = Reduce
      { forall rep. Reduce rep -> Commutativity
    redComm :: Commutativity,
        forall rep. Reduce rep -> Lambda rep
    redLambda :: Lambda rep,
        forall rep. Reduce rep -> [SubExp]
    redNeutral :: [SubExp]
      }
      deriving (Reduce rep -> Reduce rep -> Bool
    (Reduce rep -> Reduce rep -> Bool)
    -> (Reduce rep -> Reduce rep -> Bool) -> Eq (Reduce rep)
    forall rep. RepTypes rep => Reduce rep -> Reduce rep -> Bool
    forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
    $c== :: forall rep. RepTypes rep => Reduce rep -> Reduce rep -> Bool
    == :: Reduce rep -> Reduce rep -> Bool
    $c/= :: forall rep. RepTypes rep => Reduce rep -> Reduce rep -> Bool
    /= :: Reduce rep -> Reduce rep -> Bool
    Eq, Eq (Reduce rep)
    Eq (Reduce rep) =>
    (Reduce rep -> Reduce rep -> Ordering)
    -> (Reduce rep -> Reduce rep -> Bool)
    -> (Reduce rep -> Reduce rep -> Bool)
    -> (Reduce rep -> Reduce rep -> Bool)
    -> (Reduce rep -> Reduce rep -> Bool)
    -> (Reduce rep -> Reduce rep -> Reduce rep)
    -> (Reduce rep -> Reduce rep -> Reduce rep)
    -> Ord (Reduce rep)
    Reduce rep -> Reduce rep -> Bool
    Reduce rep -> Reduce rep -> Ordering
    Reduce rep -> Reduce rep -> Reduce rep
    forall a.
    Eq a =>
    (a -> a -> Ordering)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> Bool)
    -> (a -> a -> a)
    -> (a -> a -> a)
    -> Ord a
    forall rep. RepTypes rep => Eq (Reduce rep)
    forall rep. RepTypes rep => Reduce rep -> Reduce rep -> Bool
    forall rep. RepTypes rep => Reduce rep -> Reduce rep -> Ordering
    forall rep. RepTypes rep => Reduce rep -> Reduce rep -> Reduce rep
    $ccompare :: forall rep. RepTypes rep => Reduce rep -> Reduce rep -> Ordering
    compare :: Reduce rep -> Reduce rep -> Ordering
    $c< :: forall rep. RepTypes rep => Reduce rep -> Reduce rep -> Bool
    < :: Reduce rep -> Reduce rep -> Bool
    $c<= :: forall rep. RepTypes rep => Reduce rep -> Reduce rep -> Bool
    <= :: Reduce rep -> Reduce rep -> Bool
    $c> :: forall rep. RepTypes rep => Reduce rep -> Reduce rep -> Bool
    > :: Reduce rep -> Reduce rep -> Bool
    $c>= :: forall rep. RepTypes rep => Reduce rep -> Reduce rep -> Bool
    >= :: Reduce rep -> Reduce rep -> Bool
    $cmax :: forall rep. RepTypes rep => Reduce rep -> Reduce rep -> Reduce rep
    max :: Reduce rep -> Reduce rep -> Reduce rep
    $cmin :: forall rep. RepTypes rep => Reduce rep -> Reduce rep -> Reduce rep
    min :: Reduce rep -> Reduce rep -> Reduce rep
    Ord, Int -> Reduce rep -> ShowS
    [Reduce rep] -> ShowS
    Reduce rep -> String
    (Int -> Reduce rep -> ShowS)
    -> (Reduce rep -> String)
    -> ([Reduce rep] -> ShowS)
    -> Show (Reduce rep)
    forall rep. RepTypes rep => Int -> Reduce rep -> ShowS
    forall rep. RepTypes rep => [Reduce rep] -> ShowS
    forall rep. RepTypes rep => Reduce rep -> String
    forall a.
    (Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
    $cshowsPrec :: forall rep. RepTypes rep => Int -> Reduce rep -> ShowS
    showsPrec :: Int -> Reduce rep -> ShowS
    $cshow :: forall rep. RepTypes rep => Reduce rep -> String
    show :: Reduce rep -> String
    $cshowList :: forall rep. RepTypes rep => [Reduce rep] -> ShowS
    showList :: [Reduce rep] -> ShowS
    Show)

    -- | What are the sizes of reduction results produced by these 'Reduce's?
    redSizes :: [Reduce rep] -> [Int]
    redSizes :: forall rep. [Reduce rep] -> [Int]
    redSizes = (Reduce rep -> Int) -> [Reduce rep] -> [Int]
    forall a b. (a -> b) -> [a] -> [b]
    map ([SubExp] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length ([SubExp] -> Int) -> (Reduce rep -> [SubExp]) -> Reduce rep -> Int
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Reduce rep -> [SubExp]
    forall rep. Reduce rep -> [SubExp]
    redNeutral)

    -- | How many reduction results are produced by these 'Reduce's?
    redResults :: [Reduce rep] -> Int
    redResults :: forall rep. [Reduce rep] -> Int
    redResults = [Int] -> Int
    forall a. Num a => [a] -> a
    forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
    sum ([Int] -> Int) -> ([Reduce rep] -> [Int]) -> [Reduce rep] -> Int
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . [Reduce rep] -> [Int]
    forall rep. [Reduce rep] -> [Int]
    redSizes

    -- | Combine multiple reduction operators to a single operator.
    singleReduce :: (Buildable rep) => [Reduce rep] -> Reduce rep
    singleReduce :: forall rep. Buildable rep => [Reduce rep] -> Reduce rep
    singleReduce [Reduce rep]
    reds =
      let red_nes :: [SubExp]
    red_nes = (Reduce rep -> [SubExp]) -> [Reduce rep] -> [SubExp]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap Reduce rep -> [SubExp]
    forall rep. Reduce rep -> [SubExp]
    redNeutral [Reduce rep]
    reds
          red_lam :: Lambda rep
    red_lam = [Lambda rep] -> Lambda rep
    forall rep. Buildable rep => [Lambda rep] -> Lambda rep
    singleBinOp ([Lambda rep] -> Lambda rep) -> [Lambda rep] -> Lambda rep
    forall a b. (a -> b) -> a -> b
    $ (Reduce rep -> Lambda rep) -> [Reduce rep] -> [Lambda rep]
    forall a b. (a -> b) -> [a] -> [b]
    map Reduce rep -> Lambda rep
    forall rep. Reduce rep -> Lambda rep
    redLambda [Reduce rep]
    reds
       in Commutativity -> Lambda rep -> [SubExp] -> Reduce rep
    forall rep. Commutativity -> Lambda rep -> [SubExp] -> Reduce rep
    Reduce ([Commutativity] -> Commutativity
    forall a. Monoid a => [a] -> a
    mconcat ((Reduce rep -> Commutativity) -> [Reduce rep] -> [Commutativity]
    forall a b. (a -> b) -> [a] -> [b]
    map Reduce rep -> Commutativity
    forall rep. Reduce rep -> Commutativity
    redComm [Reduce rep]
    reds)) Lambda rep
    red_lam [SubExp]
    red_nes

    -- | The types produced by a single 'Screma', given the size of the
    -- input array.
    scremaType :: SubExp -> ScremaForm rep -> [Type]
    scremaType :: forall rep. SubExp -> ScremaForm rep -> [Type]
    scremaType SubExp
    w (ScremaForm Lambda rep
    map_lam [Scan rep]
    scans [Reduce rep]
    reds) =
      [Type]
    scan_tps [Type] -> [Type] -> [Type]
    forall a. [a] -> [a] -> [a]
    ++ [Type]
    red_tps [Type] -> [Type] -> [Type]
    forall a. [a] -> [a] -> [a]
    ++ (Type -> Type) -> [Type] -> [Type]
    forall a b. (a -> b) -> [a] -> [b]
    map (Type -> SubExp -> Type
    forall d.
    ArrayShape (ShapeBase d) =>
    TypeBase (ShapeBase d) NoUniqueness
    -> d -> TypeBase (ShapeBase d) NoUniqueness
    `arrayOfRow` SubExp
    w) [Type]
    map_tps
      where
        scan_tps :: [Type]
    scan_tps =
          (Type -> Type) -> [Type] -> [Type]
    forall a b. (a -> b) -> [a] -> [b]
    map (Type -> SubExp -> Type
    forall d.
    ArrayShape (ShapeBase d) =>
    TypeBase (ShapeBase d) NoUniqueness
    -> d -> TypeBase (ShapeBase d) NoUniqueness
    `arrayOfRow` SubExp
    w) ([Type] -> [Type]) -> [Type] -> [Type]
    forall a b. (a -> b) -> a -> b
    $
            (Scan rep -> [Type]) -> [Scan rep] -> [Type]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap (Lambda rep -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType (Lambda rep -> [Type])
    -> (Scan rep -> Lambda rep) -> Scan rep -> [Type]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Scan rep -> Lambda rep
    forall rep. Scan rep -> Lambda rep
    scanLambda) [Scan rep]
    scans
        red_tps :: [Type]
    red_tps = (Reduce rep -> [Type]) -> [Reduce rep] -> [Type]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap (Lambda rep -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType (Lambda rep -> [Type])
    -> (Reduce rep -> Lambda rep) -> Reduce rep -> [Type]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Reduce rep -> Lambda rep
    forall rep. Reduce rep -> Lambda rep
    redLambda) [Reduce rep]
    reds
        map_tps :: [Type]
    map_tps = Int -> [Type] -> [Type]
    forall a. Int -> [a] -> [a]
    drop ([Type] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [Type]
    scan_tps Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + [Type] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [Type]
    red_tps) ([Type] -> [Type]) -> [Type] -> [Type]
    forall a b. (a -> b) -> a -> b
    $ Lambda rep -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda rep
    map_lam

    -- | Construct a lambda that takes parameters of the given types and
    -- simply returns them unchanged.
    mkIdentityLambda ::
      (Buildable rep, MonadFreshNames m) =>
      [Type] ->
      m (Lambda rep)
    mkIdentityLambda :: forall rep (m :: * -> *).
    (Buildable rep, MonadFreshNames m) =>
    [Type] -> m (Lambda rep)
    mkIdentityLambda [Type]
    ts = do
      [Param Type]
    params <- (Type -> m (Param Type)) -> [Type] -> m [Param Type]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM (String -> Type -> m (Param Type)
    forall (m :: * -> *) dec.
    MonadFreshNames m =>
    String -> dec -> m (Param dec)
    newParam String
    "x") [Type]
    ts
      Lambda rep -> m (Lambda rep)
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure
        Lambda
          { lambdaParams :: [LParam rep]
    lambdaParams = [Param Type]
    [LParam rep]
    params,
            lambdaBody :: Body rep
    lambdaBody = Stms rep -> Result -> Body rep
    forall rep. Buildable rep => Stms rep -> Result -> Body rep
    mkBody Stms rep
    forall a. Monoid a => a
    mempty (Result -> Body rep) -> Result -> Body rep
    forall a b. (a -> b) -> a -> b
    $ [VName] -> Result
    varsRes ([VName] -> Result) -> [VName] -> Result
    forall a b. (a -> b) -> a -> b
    $ (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]
    params,
            lambdaReturnType :: [Type]
    lambdaReturnType = [Type]
    ts
          }

    -- | Is the given lambda an identity lambda?
    isIdentityLambda :: Lambda rep -> Bool
    isIdentityLambda :: forall rep. Lambda rep -> Bool
    isIdentityLambda Lambda rep
    lam =
      (SubExpRes -> SubExp) -> Result -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExpRes -> SubExp
    resSubExp (Body rep -> Result
    forall rep. Body rep -> Result
    bodyResult (Lambda rep -> Body rep
    forall rep. Lambda rep -> Body rep
    lambdaBody Lambda rep
    lam))
        [SubExp] -> [SubExp] -> Bool
    forall a. Eq a => a -> a -> Bool
    == (Param (LParamInfo rep) -> SubExp)
    -> [Param (LParamInfo rep)] -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map (VName -> SubExp
    Var (VName -> SubExp)
    -> (Param (LParamInfo rep) -> VName)
    -> Param (LParamInfo rep)
    -> SubExp
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Param (LParamInfo rep) -> VName
    forall dec. Param dec -> VName
    paramName) (Lambda rep -> [Param (LParamInfo rep)]
    forall {rep}. Lambda rep -> [Param (LParamInfo rep)]
    lambdaParams Lambda rep
    lam)

    -- | A lambda with no parameters that returns no values.
    nilFn :: (Buildable rep) => Lambda rep
    nilFn :: forall rep. Buildable rep => Lambda rep
    nilFn = [LParam rep] -> [Type] -> Body rep -> Lambda rep
    forall rep. [LParam rep] -> [Type] -> Body rep -> Lambda rep
    Lambda [Param Type]
    [LParam rep]
    forall a. Monoid a => a
    mempty [Type]
    forall a. Monoid a => a
    mempty (Stms rep -> Result -> Body rep
    forall rep. Buildable rep => Stms rep -> Result -> Body rep
    mkBody Stms rep
    forall a. Monoid a => a
    mempty Result
    forall a. Monoid a => a
    mempty)

    -- | Construct a Screma with possibly multiple scans, and
    -- the given map function.
    scanomapSOAC :: [Scan rep] -> Lambda rep -> ScremaForm rep
    scanomapSOAC :: forall rep. [Scan rep] -> Lambda rep -> ScremaForm rep
    scanomapSOAC [Scan rep]
    scans Lambda rep
    lam = Lambda rep -> [Scan rep] -> [Reduce rep] -> ScremaForm rep
    forall rep.
    Lambda rep -> [Scan rep] -> [Reduce rep] -> ScremaForm rep
    ScremaForm Lambda rep
    lam [Scan rep]
    scans []

    -- | Construct a Screma with possibly multiple reductions, and
    -- the given map function.
    redomapSOAC :: [Reduce rep] -> Lambda rep -> ScremaForm rep
    redomapSOAC :: forall rep. [Reduce rep] -> Lambda rep -> ScremaForm rep
    redomapSOAC [Reduce rep]
    reds Lambda rep
    lam = Lambda rep -> [Scan rep] -> [Reduce rep] -> ScremaForm rep
    forall rep.
    Lambda rep -> [Scan rep] -> [Reduce rep] -> ScremaForm rep
    ScremaForm Lambda rep
    lam [] [Reduce rep]
    reds

    -- | Construct a Screma with possibly multiple scans, and identity map
    -- function.
    scanSOAC ::
      (Buildable rep, MonadFreshNames m) =>
      [Scan rep] ->
      m (ScremaForm rep)
    scanSOAC :: forall rep (m :: * -> *).
    (Buildable rep, MonadFreshNames m) =>
    [Scan rep] -> m (ScremaForm rep)
    scanSOAC [Scan rep]
    scans = [Scan rep] -> Lambda rep -> ScremaForm rep
    forall rep. [Scan rep] -> Lambda rep -> ScremaForm rep
    scanomapSOAC [Scan rep]
    scans (Lambda rep -> ScremaForm rep)
    -> m (Lambda rep) -> m (ScremaForm rep)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> [Type] -> m (Lambda rep)
    forall rep (m :: * -> *).
    (Buildable rep, MonadFreshNames m) =>
    [Type] -> m (Lambda rep)
    mkIdentityLambda [Type]
    ts
      where
        ts :: [Type]
    ts = (Scan rep -> [Type]) -> [Scan rep] -> [Type]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap (Lambda rep -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType (Lambda rep -> [Type])
    -> (Scan rep -> Lambda rep) -> Scan rep -> [Type]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Scan rep -> Lambda rep
    forall rep. Scan rep -> Lambda rep
    scanLambda) [Scan rep]
    scans

    -- | Construct a Screma with possibly multiple reductions, and
    -- identity map function.
    reduceSOAC ::
      (Buildable rep, MonadFreshNames m) =>
      [Reduce rep] ->
      m (ScremaForm rep)
    reduceSOAC :: forall rep (m :: * -> *).
    (Buildable rep, MonadFreshNames m) =>
    [Reduce rep] -> m (ScremaForm rep)
    reduceSOAC [Reduce rep]
    reds = [Reduce rep] -> Lambda rep -> ScremaForm rep
    forall rep. [Reduce rep] -> Lambda rep -> ScremaForm rep
    redomapSOAC [Reduce rep]
    reds (Lambda rep -> ScremaForm rep)
    -> m (Lambda rep) -> m (ScremaForm rep)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> [Type] -> m (Lambda rep)
    forall rep (m :: * -> *).
    (Buildable rep, MonadFreshNames m) =>
    [Type] -> m (Lambda rep)
    mkIdentityLambda [Type]
    ts
      where
        ts :: [Type]
    ts = (Reduce rep -> [Type]) -> [Reduce rep] -> [Type]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap (Lambda rep -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType (Lambda rep -> [Type])
    -> (Reduce rep -> Lambda rep) -> Reduce rep -> [Type]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Reduce rep -> Lambda rep
    forall rep. Reduce rep -> Lambda rep
    redLambda) [Reduce rep]
    reds

    -- | Construct a Screma corresponding to a map.
    mapSOAC :: Lambda rep -> ScremaForm rep
    mapSOAC :: forall rep. Lambda rep -> ScremaForm rep
    mapSOAC Lambda rep
    lam = Lambda rep -> [Scan rep] -> [Reduce rep] -> ScremaForm rep
    forall rep.
    Lambda rep -> [Scan rep] -> [Reduce rep] -> ScremaForm rep
    ScremaForm Lambda rep
    lam [] []

    -- | Does this Screma correspond to a scan-map composition?
    isScanomapSOAC :: ScremaForm rep -> Maybe ([Scan rep], Lambda rep)
    isScanomapSOAC :: forall rep. ScremaForm rep -> Maybe ([Scan rep], Lambda rep)
    isScanomapSOAC (ScremaForm Lambda rep
    map_lam [Scan rep]
    scans [Reduce rep]
    reds) = do
      Bool -> Maybe ()
    forall (f :: * -> *). Alternative f => Bool -> f ()
    guard (Bool -> Maybe ()) -> Bool -> Maybe ()
    forall a b. (a -> b) -> a -> b
    $ [Reduce rep] -> Bool
    forall a. [a] -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null [Reduce rep]
    reds
      Bool -> Maybe ()
    forall (f :: * -> *). Alternative f => Bool -> f ()
    guard (Bool -> Maybe ()) -> Bool -> Maybe ()
    forall a b. (a -> b) -> a -> b
    $ Bool -> Bool
    not (Bool -> Bool) -> Bool -> Bool
    forall a b. (a -> b) -> a -> b
    $ [Scan rep] -> Bool
    forall a. [a] -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null [Scan rep]
    scans
      ([Scan rep], Lambda rep) -> Maybe ([Scan rep], Lambda rep)
    forall a. a -> Maybe a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ([Scan rep]
    scans, Lambda rep
    map_lam)

    -- | Does this Screma correspond to pure scan?
    isScanSOAC :: ScremaForm rep -> Maybe [Scan rep]
    isScanSOAC :: forall rep. ScremaForm rep -> Maybe [Scan rep]
    isScanSOAC ScremaForm rep
    form = do
      ([Scan rep]
    scans, Lambda rep
    map_lam) <- ScremaForm rep -> Maybe ([Scan rep], Lambda rep)
    forall rep. ScremaForm rep -> Maybe ([Scan rep], Lambda rep)
    isScanomapSOAC ScremaForm rep
    form
      Bool -> Maybe ()
    forall (f :: * -> *). Alternative f => Bool -> f ()
    guard (Bool -> Maybe ()) -> Bool -> Maybe ()
    forall a b. (a -> b) -> a -> b
    $ Lambda rep -> Bool
    forall rep. Lambda rep -> Bool
    isIdentityLambda Lambda rep
    map_lam
      [Scan rep] -> Maybe [Scan rep]
    forall a. a -> Maybe a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [Scan rep]
    scans

    -- | Does this Screma correspond to a reduce-map composition?
    isRedomapSOAC :: ScremaForm rep -> Maybe ([Reduce rep], Lambda rep)
    isRedomapSOAC :: forall rep. ScremaForm rep -> Maybe ([Reduce rep], Lambda rep)
    isRedomapSOAC (ScremaForm Lambda rep
    map_lam [Scan rep]
    scans [Reduce rep]
    reds) = do
      Bool -> Maybe ()
    forall (f :: * -> *). Alternative f => Bool -> f ()
    guard (Bool -> Maybe ()) -> Bool -> Maybe ()
    forall a b. (a -> b) -> a -> b
    $ [Scan rep] -> Bool
    forall a. [a] -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null [Scan rep]
    scans
      Bool -> Maybe ()
    forall (f :: * -> *). Alternative f => Bool -> f ()
    guard (Bool -> Maybe ()) -> Bool -> Maybe ()
    forall a b. (a -> b) -> a -> b
    $ Bool -> Bool
    not (Bool -> Bool) -> Bool -> Bool
    forall a b. (a -> b) -> a -> b
    $ [Reduce rep] -> Bool
    forall a. [a] -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null [Reduce rep]
    reds
      ([Reduce rep], Lambda rep) -> Maybe ([Reduce rep], Lambda rep)
    forall a. a -> Maybe a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ([Reduce rep]
    reds, Lambda rep
    map_lam)

    -- | Does this Screma correspond to a pure reduce?
    isReduceSOAC :: ScremaForm rep -> Maybe [Reduce rep]
    isReduceSOAC :: forall rep. ScremaForm rep -> Maybe [Reduce rep]
    isReduceSOAC ScremaForm rep
    form = do
      ([Reduce rep]
    reds, Lambda rep
    map_lam) <- ScremaForm rep -> Maybe ([Reduce rep], Lambda rep)
    forall rep. ScremaForm rep -> Maybe ([Reduce rep], Lambda rep)
    isRedomapSOAC ScremaForm rep
    form
      Bool -> Maybe ()
    forall (f :: * -> *). Alternative f => Bool -> f ()
    guard (Bool -> Maybe ()) -> Bool -> Maybe ()
    forall a b. (a -> b) -> a -> b
    $ Lambda rep -> Bool
    forall rep. Lambda rep -> Bool
    isIdentityLambda Lambda rep
    map_lam
      [Reduce rep] -> Maybe [Reduce rep]
    forall a. a -> Maybe a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [Reduce rep]
    reds

    -- | Does this Screma correspond to a simple map, without any
    -- reduction or scan results?
    isMapSOAC :: ScremaForm rep -> Maybe (Lambda rep)
    isMapSOAC :: forall rep. ScremaForm rep -> Maybe (Lambda rep)
    isMapSOAC (ScremaForm Lambda rep
    map_lam [Scan rep]
    scans [Reduce rep]
    reds) = do
      Bool -> Maybe ()
    forall (f :: * -> *). Alternative f => Bool -> f ()
    guard (Bool -> Maybe ()) -> Bool -> Maybe ()
    forall a b. (a -> b) -> a -> b
    $ [Scan rep] -> Bool
    forall a. [a] -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null [Scan rep]
    scans
      Bool -> Maybe ()
    forall (f :: * -> *). Alternative f => Bool -> f ()
    guard (Bool -> Maybe ()) -> Bool -> Maybe ()
    forall a b. (a -> b) -> a -> b
    $ [Reduce rep] -> Bool
    forall a. [a] -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null [Reduce rep]
    reds
      Lambda rep -> Maybe (Lambda rep)
    forall a. a -> Maybe a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Lambda rep
    map_lam

    -- | @splitScatterResults <scatter specification> <results>@
    --
    -- Splits the results array into indices and values according to the
    -- specification.
    --
    -- See 'groupScatterResults' for more information.
    splitScatterResults :: [(Shape, Int, array)] -> [a] -> ([a], [a])
    splitScatterResults :: forall array a. [(Shape, Int, array)] -> [a] -> ([a], [a])
    splitScatterResults [(Shape, Int, array)]
    output_spec [a]
    results =
      let ([Shape]
    shapes, [Int]
    ns, [array]
    _) = [(Shape, Int, array)] -> ([Shape], [Int], [array])
    forall a b c. [(a, b, c)] -> ([a], [b], [c])
    unzip3 [(Shape, Int, array)]
    output_spec
          num_indices :: Int
    num_indices = [Int] -> Int
    forall a. Num a => [a] -> a
    forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
    sum ([Int] -> Int) -> [Int] -> Int
    forall a b. (a -> b) -> a -> b
    $ (Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith Int -> Int -> Int
    forall a. Num a => a -> a -> a
    (*) [Int]
    ns ([Int] -> [Int]) -> [Int] -> [Int]
    forall a b. (a -> b) -> a -> b
    $ (Shape -> Int) -> [Shape] -> [Int]
    forall a b. (a -> b) -> [a] -> [b]
    map Shape -> Int
    forall a. ShapeBase a -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [Shape]
    shapes
       in Int -> [a] -> ([a], [a])
    forall a. Int -> [a] -> ([a], [a])
    splitAt Int
    num_indices [a]
    results

    -- | @groupScatterResults' <scatter specification> <results>@
    --
    -- Blocks the index values and result values of <results> according to
    -- the specification. This is the simpler version of
    -- @groupScatterResults@, which doesn't return any information about
    -- shapes or output arrays.
    --
    -- See 'groupScatterResults' for more information,
    groupScatterResults' :: [(Shape, Int, array)] -> [a] -> [([a], a)]
    groupScatterResults' :: forall array a. [(Shape, Int, array)] -> [a] -> [([a], a)]
    groupScatterResults' [(Shape, Int, array)]
    output_spec [a]
    results =
      let ([a]
    indices, [a]
    values) = [(Shape, Int, array)] -> [a] -> ([a], [a])
    forall array a. [(Shape, Int, array)] -> [a] -> ([a], [a])
    splitScatterResults [(Shape, Int, array)]
    output_spec [a]
    results
          ([Shape]
    shapes, [Int]
    ns, [array]
    _) = [(Shape, Int, array)] -> ([Shape], [Int], [array])
    forall a b c. [(a, b, c)] -> ([a], [b], [c])
    unzip3 [(Shape, Int, array)]
    output_spec
          chunk_sizes :: [Int]
    chunk_sizes =
            [[Int]] -> [Int]
    forall (t :: * -> *) a. Foldable t => t [a] -> [a]
    concat ([[Int]] -> [Int]) -> [[Int]] -> [Int]
    forall a b. (a -> b) -> a -> b
    $ (Shape -> Int -> [Int]) -> [Shape] -> [Int] -> [[Int]]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith (\Shape
    shp Int
    n -> Int -> Int -> [Int]
    forall a. Int -> a -> [a]
    replicate Int
    n (Int -> [Int]) -> Int -> [Int]
    forall a b. (a -> b) -> a -> b
    $ Shape -> Int
    forall a. ShapeBase a -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length Shape
    shp) [Shape]
    shapes [Int]
    ns
       in [[a]] -> [a] -> [([a], a)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip ([Int] -> [a] -> [[a]]
    forall a. [Int] -> [a] -> [[a]]
    chunks [Int]
    chunk_sizes [a]
    indices) [a]
    values

    -- | @groupScatterResults <scatter specification> <results>@
    --
    -- Blocks the index values and result values of <results> according to the
    -- <scatter specification>.
    --
    -- This function is used for extracting and grouping the results of a
    -- scatter. In the SOACS representation, the lambda inside a 'Scatter' returns
    -- all indices and values as one big list. This function groups each value with
    -- its corresponding indices (as determined by the t'Shape' of the output array).
    --
    -- The elements of the resulting list correspond to the shape and name of the
    -- output parameters, in addition to a list of values written to that output
    -- parameter, along with the array indices marking where to write them to.
    --
    -- See 'Scatter' for more information.
    groupScatterResults :: ScatterSpec array -> [a] -> [(Shape, array, [([a], a)])]
    groupScatterResults :: forall array a.
    ScatterSpec array -> [a] -> [(Shape, array, [([a], a)])]
    groupScatterResults ScatterSpec array
    output_spec [a]
    results =
      let ([Shape]
    shapes, [Int]
    ns, [array]
    arrays) = ScatterSpec array -> ([Shape], [Int], [array])
    forall a b c. [(a, b, c)] -> ([a], [b], [c])
    unzip3 ScatterSpec array
    output_spec
       in ScatterSpec array -> [a] -> [([a], a)]
    forall array a. [(Shape, Int, array)] -> [a] -> [([a], a)]
    groupScatterResults' ScatterSpec array
    output_spec [a]
    results
            [([a], a)] -> ([([a], a)] -> [[([a], a)]]) -> [[([a], a)]]
    forall a b. a -> (a -> b) -> b
    & [Int] -> [([a], a)] -> [[([a], a)]]
    forall a. [Int] -> [a] -> [[a]]
    chunks [Int]
    ns
            [[([a], a)]]
    -> ([[([a], a)]] -> [(Shape, array, [([a], a)])])
    -> [(Shape, array, [([a], a)])]
    forall a b. a -> (a -> b) -> b
    & [Shape] -> [array] -> [[([a], a)]] -> [(Shape, array, [([a], a)])]
    forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
    zip3 [Shape]
    shapes [array]
    arrays

    -- | Like 'Mapper', but just for 'SOAC's.
    data SOACMapper frep trep m = SOACMapper
      { forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> SubExp -> m SubExp
    mapOnSOACSubExp :: SubExp -> m SubExp,
        forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
    mapOnSOACLambda :: Lambda frep -> m (Lambda trep),
        forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> VName -> m VName
    mapOnSOACVName :: VName -> m VName
      }

    -- | A mapper that simply returns the SOAC verbatim.
    identitySOACMapper :: forall rep m. (Monad m) => SOACMapper rep rep m
    identitySOACMapper :: forall rep (m :: * -> *). Monad m => SOACMapper rep rep m
    identitySOACMapper =
      SOACMapper
        { mapOnSOACSubExp :: SubExp -> m SubExp
    mapOnSOACSubExp = SubExp -> m SubExp
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure,
          mapOnSOACLambda :: Lambda rep -> m (Lambda rep)
    mapOnSOACLambda = Lambda rep -> m (Lambda rep)
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure,
          mapOnSOACVName :: VName -> m VName
    mapOnSOACVName = VName -> m VName
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure
        }

    -- | Map a monadic action across the immediate children of a
    -- SOAC.  The mapping does not descend recursively into subexpressions
    -- and is done left-to-right.
    mapSOACM ::
      (Monad m) =>
      SOACMapper frep trep m ->
      SOAC frep ->
      m (SOAC trep)
    mapSOACM :: forall (m :: * -> *) frep trep.
    Monad m =>
    SOACMapper frep trep m -> SOAC frep -> m (SOAC trep)
    mapSOACM SOACMapper frep trep m
    tv (JVP [SubExp]
    args [SubExp]
    vec Lambda frep
    lam) =
      [SubExp] -> [SubExp] -> Lambda trep -> SOAC trep
    forall rep. [SubExp] -> [SubExp] -> Lambda rep -> SOAC rep
    JVP
        ([SubExp] -> [SubExp] -> Lambda trep -> SOAC trep)
    -> m [SubExp] -> m ([SubExp] -> Lambda trep -> SOAC trep)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (SubExp -> m SubExp) -> [SubExp] -> m [SubExp]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM (SOACMapper frep trep m -> SubExp -> m SubExp
    forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> SubExp -> m SubExp
    mapOnSOACSubExp SOACMapper frep trep m
    tv) [SubExp]
    args
        m ([SubExp] -> Lambda trep -> SOAC trep)
    -> m [SubExp] -> m (Lambda trep -> SOAC trep)
    forall a b. m (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> (SubExp -> m SubExp) -> [SubExp] -> m [SubExp]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM (SOACMapper frep trep m -> SubExp -> m SubExp
    forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> SubExp -> m SubExp
    mapOnSOACSubExp SOACMapper frep trep m
    tv) [SubExp]
    vec
        m (Lambda trep -> SOAC trep) -> m (Lambda trep) -> m (SOAC trep)
    forall a b. m (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
    forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
    mapOnSOACLambda SOACMapper frep trep m
    tv Lambda frep
    lam
    mapSOACM SOACMapper frep trep m
    tv (VJP [SubExp]
    args [SubExp]
    vec Lambda frep
    lam) =
      [SubExp] -> [SubExp] -> Lambda trep -> SOAC trep
    forall rep. [SubExp] -> [SubExp] -> Lambda rep -> SOAC rep
    VJP
        ([SubExp] -> [SubExp] -> Lambda trep -> SOAC trep)
    -> m [SubExp] -> m ([SubExp] -> Lambda trep -> SOAC trep)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (SubExp -> m SubExp) -> [SubExp] -> m [SubExp]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM (SOACMapper frep trep m -> SubExp -> m SubExp
    forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> SubExp -> m SubExp
    mapOnSOACSubExp SOACMapper frep trep m
    tv) [SubExp]
    args
        m ([SubExp] -> Lambda trep -> SOAC trep)
    -> m [SubExp] -> m (Lambda trep -> SOAC trep)
    forall a b. m (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> (SubExp -> m SubExp) -> [SubExp] -> m [SubExp]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM (SOACMapper frep trep m -> SubExp -> m SubExp
    forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> SubExp -> m SubExp
    mapOnSOACSubExp SOACMapper frep trep m
    tv) [SubExp]
    vec
        m (Lambda trep -> SOAC trep) -> m (Lambda trep) -> m (SOAC trep)
    forall a b. m (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
    forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
    mapOnSOACLambda SOACMapper frep trep m
    tv Lambda frep
    lam
    mapSOACM SOACMapper frep trep m
    tv (Stream SubExp
    size [VName]
    arrs [SubExp]
    accs Lambda frep
    lam) =
      SubExp -> [VName] -> [SubExp] -> Lambda trep -> SOAC trep
    forall rep. SubExp -> [VName] -> [SubExp] -> Lambda rep -> SOAC rep
    Stream
        (SubExp -> [VName] -> [SubExp] -> Lambda trep -> SOAC trep)
    -> m SubExp -> m ([VName] -> [SubExp] -> Lambda trep -> SOAC trep)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> SOACMapper frep trep m -> SubExp -> m SubExp
    forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> SubExp -> m SubExp
    mapOnSOACSubExp SOACMapper frep trep m
    tv SubExp
    size
        m ([VName] -> [SubExp] -> Lambda trep -> SOAC trep)
    -> m [VName] -> m ([SubExp] -> Lambda trep -> SOAC trep)
    forall a b. m (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> (VName -> m VName) -> [VName] -> m [VName]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM (SOACMapper frep trep m -> VName -> m VName
    forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> VName -> m VName
    mapOnSOACVName SOACMapper frep trep m
    tv) [VName]
    arrs
        m ([SubExp] -> Lambda trep -> SOAC trep)
    -> m [SubExp] -> m (Lambda trep -> SOAC trep)
    forall a b. m (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> (SubExp -> m SubExp) -> [SubExp] -> m [SubExp]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM (SOACMapper frep trep m -> SubExp -> m SubExp
    forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> SubExp -> m SubExp
    mapOnSOACSubExp SOACMapper frep trep m
    tv) [SubExp]
    accs
        m (Lambda trep -> SOAC trep) -> m (Lambda trep) -> m (SOAC trep)
    forall a b. m (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
    forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
    mapOnSOACLambda SOACMapper frep trep m
    tv Lambda frep
    lam
    mapSOACM SOACMapper frep trep m
    tv (Scatter SubExp
    w [VName]
    ivs ScatterSpec VName
    as Lambda frep
    lam) =
      SubExp -> [VName] -> ScatterSpec VName -> Lambda trep -> SOAC trep
    forall rep.
    SubExp -> [VName] -> ScatterSpec VName -> Lambda rep -> SOAC rep
    Scatter
        (SubExp
     -> [VName] -> ScatterSpec VName -> Lambda trep -> SOAC trep)
    -> m SubExp
    -> m ([VName] -> ScatterSpec VName -> Lambda trep -> SOAC trep)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> SOACMapper frep trep m -> SubExp -> m SubExp
    forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> SubExp -> m SubExp
    mapOnSOACSubExp SOACMapper frep trep m
    tv SubExp
    w
        m ([VName] -> ScatterSpec VName -> Lambda trep -> SOAC trep)
    -> m [VName] -> m (ScatterSpec VName -> Lambda trep -> SOAC trep)
    forall a b. m (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> (VName -> m VName) -> [VName] -> m [VName]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM (SOACMapper frep trep m -> VName -> m VName
    forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> VName -> m VName
    mapOnSOACVName SOACMapper frep trep m
    tv) [VName]
    ivs
        m (ScatterSpec VName -> Lambda trep -> SOAC trep)
    -> m (ScatterSpec VName) -> m (Lambda trep -> SOAC trep)
    forall a b. m (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> ((Shape, Int, VName) -> m (Shape, Int, VName))
    -> ScatterSpec VName -> m (ScatterSpec VName)
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM
          ( \(Shape
    aw, Int
    an, VName
    a) ->
              (,,)
                (Shape -> Int -> VName -> (Shape, Int, VName))
    -> m Shape -> m (Int -> VName -> (Shape, Int, VName))
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (SubExp -> m SubExp) -> Shape -> m Shape
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> ShapeBase a -> m (ShapeBase b)
    mapM (SOACMapper frep trep m -> SubExp -> m SubExp
    forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> SubExp -> m SubExp
    mapOnSOACSubExp SOACMapper frep trep m
    tv) Shape
    aw
                m (Int -> VName -> (Shape, Int, VName))
    -> m Int -> m (VName -> (Shape, Int, VName))
    forall a b. m (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> Int -> m Int
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Int
    an
                m (VName -> (Shape, Int, VName))
    -> m VName -> m (Shape, Int, VName)
    forall a b. m (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> SOACMapper frep trep m -> VName -> m VName
    forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> VName -> m VName
    mapOnSOACVName SOACMapper frep trep m
    tv VName
    a
          )
          ScatterSpec VName
    as
        m (Lambda trep -> SOAC trep) -> m (Lambda trep) -> m (SOAC trep)
    forall a b. m (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
    forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
    mapOnSOACLambda SOACMapper frep trep m
    tv Lambda frep
    lam
    mapSOACM SOACMapper frep trep m
    tv (Hist SubExp
    w [VName]
    arrs [HistOp frep]
    ops Lambda frep
    bucket_fun) =
      SubExp -> [VName] -> [HistOp trep] -> Lambda trep -> SOAC trep
    forall rep.
    SubExp -> [VName] -> [HistOp rep] -> Lambda rep -> SOAC rep
    Hist
        (SubExp -> [VName] -> [HistOp trep] -> Lambda trep -> SOAC trep)
    -> m SubExp
    -> m ([VName] -> [HistOp trep] -> Lambda trep -> SOAC trep)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> SOACMapper frep trep m -> SubExp -> m SubExp
    forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> SubExp -> m SubExp
    mapOnSOACSubExp SOACMapper frep trep m
    tv SubExp
    w
        m ([VName] -> [HistOp trep] -> Lambda trep -> SOAC trep)
    -> m [VName] -> m ([HistOp trep] -> Lambda trep -> SOAC trep)
    forall a b. m (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> (VName -> m VName) -> [VName] -> m [VName]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM (SOACMapper frep trep m -> VName -> m VName
    forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> VName -> m VName
    mapOnSOACVName SOACMapper frep trep m
    tv) [VName]
    arrs
        m ([HistOp trep] -> Lambda trep -> SOAC trep)
    -> m [HistOp trep] -> m (Lambda trep -> SOAC trep)
    forall a b. m (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> (HistOp frep -> m (HistOp trep))
    -> [HistOp frep] -> m [HistOp trep]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM
          ( \(HistOp Shape
    shape SubExp
    rf [VName]
    op_arrs [SubExp]
    nes Lambda frep
    op) ->
              Shape
    -> SubExp -> [VName] -> [SubExp] -> Lambda trep -> HistOp trep
    forall rep.
    Shape -> SubExp -> [VName] -> [SubExp] -> Lambda rep -> HistOp rep
    HistOp
                (Shape
     -> SubExp -> [VName] -> [SubExp] -> Lambda trep -> HistOp trep)
    -> m Shape
    -> m (SubExp -> [VName] -> [SubExp] -> Lambda trep -> HistOp trep)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (SubExp -> m SubExp) -> Shape -> m Shape
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> ShapeBase a -> m (ShapeBase b)
    mapM (SOACMapper frep trep m -> SubExp -> m SubExp
    forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> SubExp -> m SubExp
    mapOnSOACSubExp SOACMapper frep trep m
    tv) Shape
    shape
                m (SubExp -> [VName] -> [SubExp] -> Lambda trep -> HistOp trep)
    -> m SubExp
    -> m ([VName] -> [SubExp] -> Lambda trep -> HistOp trep)
    forall a b. m (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> SOACMapper frep trep m -> SubExp -> m SubExp
    forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> SubExp -> m SubExp
    mapOnSOACSubExp SOACMapper frep trep m
    tv SubExp
    rf
                m ([VName] -> [SubExp] -> Lambda trep -> HistOp trep)
    -> m [VName] -> m ([SubExp] -> Lambda trep -> HistOp trep)
    forall a b. m (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> (VName -> m VName) -> [VName] -> m [VName]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM (SOACMapper frep trep m -> VName -> m VName
    forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> VName -> m VName
    mapOnSOACVName SOACMapper frep trep m
    tv) [VName]
    op_arrs
                m ([SubExp] -> Lambda trep -> HistOp trep)
    -> m [SubExp] -> m (Lambda trep -> HistOp trep)
    forall a b. m (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> (SubExp -> m SubExp) -> [SubExp] -> m [SubExp]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM (SOACMapper frep trep m -> SubExp -> m SubExp
    forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> SubExp -> m SubExp
    mapOnSOACSubExp SOACMapper frep trep m
    tv) [SubExp]
    nes
                m (Lambda trep -> HistOp trep)
    -> m (Lambda trep) -> m (HistOp trep)
    forall a b. m (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
    forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
    mapOnSOACLambda SOACMapper frep trep m
    tv Lambda frep
    op
          )
          [HistOp frep]
    ops
        m (Lambda trep -> SOAC trep) -> m (Lambda trep) -> m (SOAC trep)
    forall a b. m (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
    forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
    mapOnSOACLambda SOACMapper frep trep m
    tv Lambda frep
    bucket_fun
    mapSOACM SOACMapper frep trep m
    tv (Screma SubExp
    w [VName]
    arrs (ScremaForm Lambda frep
    map_lam [Scan frep]
    scans [Reduce frep]
    reds)) =
      SubExp -> [VName] -> ScremaForm trep -> SOAC trep
    forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
    Screma
        (SubExp -> [VName] -> ScremaForm trep -> SOAC trep)
    -> m SubExp -> m ([VName] -> ScremaForm trep -> SOAC trep)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> SOACMapper frep trep m -> SubExp -> m SubExp
    forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> SubExp -> m SubExp
    mapOnSOACSubExp SOACMapper frep trep m
    tv SubExp
    w
        m ([VName] -> ScremaForm trep -> SOAC trep)
    -> m [VName] -> m (ScremaForm trep -> SOAC trep)
    forall a b. m (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> (VName -> m VName) -> [VName] -> m [VName]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM (SOACMapper frep trep m -> VName -> m VName
    forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> VName -> m VName
    mapOnSOACVName SOACMapper frep trep m
    tv) [VName]
    arrs
        m (ScremaForm trep -> SOAC trep)
    -> m (ScremaForm trep) -> m (SOAC trep)
    forall a b. m (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> ( Lambda trep -> [Scan trep] -> [Reduce trep] -> ScremaForm trep
    forall rep.
    Lambda rep -> [Scan rep] -> [Reduce rep] -> ScremaForm rep
    ScremaForm
                (Lambda trep -> [Scan trep] -> [Reduce trep] -> ScremaForm trep)
    -> m (Lambda trep)
    -> m ([Scan trep] -> [Reduce trep] -> ScremaForm trep)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
    forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
    mapOnSOACLambda SOACMapper frep trep m
    tv Lambda frep
    map_lam
                m ([Scan trep] -> [Reduce trep] -> ScremaForm trep)
    -> m [Scan trep] -> m ([Reduce trep] -> ScremaForm trep)
    forall a b. m (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> [Scan frep] -> (Scan frep -> m (Scan trep)) -> m [Scan trep]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM
                  [Scan frep]
    scans
                  ( \(Scan Lambda frep
    red_lam [SubExp]
    red_nes) ->
                      Lambda trep -> [SubExp] -> Scan trep
    forall rep. Lambda rep -> [SubExp] -> Scan rep
    Scan
                        (Lambda trep -> [SubExp] -> Scan trep)
    -> m (Lambda trep) -> m ([SubExp] -> Scan trep)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
    forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
    mapOnSOACLambda SOACMapper frep trep m
    tv Lambda frep
    red_lam
                        m ([SubExp] -> Scan trep) -> m [SubExp] -> m (Scan trep)
    forall a b. m (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> (SubExp -> m SubExp) -> [SubExp] -> m [SubExp]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM (SOACMapper frep trep m -> SubExp -> m SubExp
    forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> SubExp -> m SubExp
    mapOnSOACSubExp SOACMapper frep trep m
    tv) [SubExp]
    red_nes
                  )
                m ([Reduce trep] -> ScremaForm trep)
    -> m [Reduce trep] -> m (ScremaForm trep)
    forall a b. m (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> [Reduce frep]
    -> (Reduce frep -> m (Reduce trep)) -> m [Reduce trep]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM
                  [Reduce frep]
    reds
                  ( \(Reduce Commutativity
    comm Lambda frep
    red_lam [SubExp]
    red_nes) ->
                      Commutativity -> Lambda trep -> [SubExp] -> Reduce trep
    forall rep. Commutativity -> Lambda rep -> [SubExp] -> Reduce rep
    Reduce Commutativity
    comm
                        (Lambda trep -> [SubExp] -> Reduce trep)
    -> m (Lambda trep) -> m ([SubExp] -> Reduce trep)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
    forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> Lambda frep -> m (Lambda trep)
    mapOnSOACLambda SOACMapper frep trep m
    tv Lambda frep
    red_lam
                        m ([SubExp] -> Reduce trep) -> m [SubExp] -> m (Reduce trep)
    forall a b. m (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> (SubExp -> m SubExp) -> [SubExp] -> m [SubExp]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM (SOACMapper frep trep m -> SubExp -> m SubExp
    forall frep trep (m :: * -> *).
    SOACMapper frep trep m -> SubExp -> m SubExp
    mapOnSOACSubExp SOACMapper frep trep m
    tv) [SubExp]
    red_nes
                  )
            )

    -- | A helper for defining 'TraverseOpStms'.
    traverseSOACStms :: (Monad m) => OpStmsTraverser m (SOAC rep) rep
    traverseSOACStms :: forall (m :: * -> *) rep.
    Monad m =>
    OpStmsTraverser m (SOAC rep) rep
    traverseSOACStms Scope rep -> Stms rep -> m (Stms rep)
    f = SOACMapper rep rep m -> SOAC rep -> m (SOAC rep)
    forall (m :: * -> *) frep trep.
    Monad m =>
    SOACMapper frep trep m -> SOAC frep -> m (SOAC trep)
    mapSOACM SOACMapper rep rep m
    mapper
      where
        mapper :: SOACMapper rep rep m
    mapper = SOACMapper Any Any m
    forall rep (m :: * -> *). Monad m => SOACMapper rep rep m
    identitySOACMapper {mapOnSOACLambda = traverseLambdaStms f}

    instance (ASTRep rep) => FreeIn (Scan rep) where
      freeIn' :: Scan rep -> FV
    freeIn' (Scan Lambda rep
    lam [SubExp]
    ne) = Lambda rep -> FV
    forall a. FreeIn a => a -> FV
    freeIn' Lambda rep
    lam FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> [SubExp] -> FV
    forall a. FreeIn a => a -> FV
    freeIn' [SubExp]
    ne

    instance (ASTRep rep) => FreeIn (Reduce rep) where
      freeIn' :: Reduce rep -> FV
    freeIn' (Reduce Commutativity
    _ Lambda rep
    lam [SubExp]
    ne) = Lambda rep -> FV
    forall a. FreeIn a => a -> FV
    freeIn' Lambda rep
    lam FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> [SubExp] -> FV
    forall a. FreeIn a => a -> FV
    freeIn' [SubExp]
    ne

    instance (ASTRep rep) => FreeIn (ScremaForm rep) where
      freeIn' :: ScremaForm rep -> FV
    freeIn' (ScremaForm Lambda rep
    scans [Scan rep]
    reds [Reduce rep]
    lam) =
        Lambda rep -> FV
    forall a. FreeIn a => a -> FV
    freeIn' Lambda rep
    scans FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> [Scan rep] -> FV
    forall a. FreeIn a => a -> FV
    freeIn' [Scan rep]
    reds FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> [Reduce rep] -> FV
    forall a. FreeIn a => a -> FV
    freeIn' [Reduce rep]
    lam

    instance (ASTRep rep) => FreeIn (HistOp rep) where
      freeIn' :: HistOp rep -> FV
    freeIn' (HistOp Shape
    w SubExp
    rf [VName]
    dests [SubExp]
    nes Lambda rep
    lam) =
        Shape -> FV
    forall a. FreeIn a => a -> FV
    freeIn' Shape
    w FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> SubExp -> FV
    forall a. FreeIn a => a -> FV
    freeIn' SubExp
    rf FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> [VName] -> FV
    forall a. FreeIn a => a -> FV
    freeIn' [VName]
    dests FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> [SubExp] -> FV
    forall a. FreeIn a => a -> FV
    freeIn' [SubExp]
    nes FV -> FV -> FV
    forall a. Semigroup a => a -> a -> a
    <> Lambda rep -> FV
    forall a. FreeIn a => a -> FV
    freeIn' Lambda rep
    lam

    instance (ASTRep rep) => FreeIn (SOAC rep) where
      freeIn' :: SOAC rep -> FV
    freeIn' = (State FV (SOAC rep) -> FV -> FV)
    -> FV -> State FV (SOAC rep) -> FV
    forall a b c. (a -> b -> c) -> b -> a -> c
    flip State FV (SOAC rep) -> FV -> FV
    forall s a. State s a -> s -> s
    execState FV
    forall a. Monoid a => a
    mempty (State FV (SOAC rep) -> FV)
    -> (SOAC rep -> State FV (SOAC rep)) -> SOAC rep -> FV
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . SOACMapper rep rep (StateT FV Identity)
    -> SOAC rep -> State FV (SOAC rep)
    forall (m :: * -> *) frep trep.
    Monad m =>
    SOACMapper frep trep m -> SOAC frep -> m (SOAC trep)
    mapSOACM SOACMapper rep rep (StateT FV Identity)
    free
        where
          walk :: (b -> s) -> b -> m b
    walk b -> s
    f b
    x = (s -> s) -> m ()
    forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
    modify (s -> s -> s
    forall a. Semigroup a => a -> a -> a
    <> b -> s
    f b
    x) m () -> m b -> m b
    forall a b. m a -> m b -> m b
    forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
    >> b -> m b
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure b
    x
          free :: SOACMapper rep rep (StateT FV Identity)
    free =
            SOACMapper
              { mapOnSOACSubExp :: SubExp -> StateT FV Identity SubExp
    mapOnSOACSubExp = (SubExp -> FV) -> SubExp -> StateT FV Identity SubExp
    forall {m :: * -> *} {s} {b}.
    (MonadState s m, Semigroup s) =>
    (b -> s) -> b -> m b
    walk SubExp -> FV
    forall a. FreeIn a => a -> FV
    freeIn',
                mapOnSOACLambda :: Lambda rep -> StateT FV Identity (Lambda rep)
    mapOnSOACLambda = (Lambda rep -> FV) -> Lambda rep -> StateT FV Identity (Lambda rep)
    forall {m :: * -> *} {s} {b}.
    (MonadState s m, Semigroup s) =>
    (b -> s) -> b -> m b
    walk Lambda rep -> FV
    forall a. FreeIn a => a -> FV
    freeIn',
                mapOnSOACVName :: VName -> StateT FV Identity VName
    mapOnSOACVName = (VName -> FV) -> VName -> StateT FV Identity VName
    forall {m :: * -> *} {s} {b}.
    (MonadState s m, Semigroup s) =>
    (b -> s) -> b -> m b
    walk VName -> FV
    forall a. FreeIn a => a -> FV
    freeIn'
              }

    instance (ASTRep rep) => Substitute (SOAC rep) where
      substituteNames :: Map VName VName -> SOAC rep -> SOAC rep
    substituteNames Map VName VName
    subst =
        Identity (SOAC rep) -> SOAC rep
    forall a. Identity a -> a
    runIdentity (Identity (SOAC rep) -> SOAC rep)
    -> (SOAC rep -> Identity (SOAC rep)) -> SOAC rep -> SOAC rep
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . SOACMapper rep rep Identity -> SOAC rep -> Identity (SOAC rep)
    forall (m :: * -> *) frep trep.
    Monad m =>
    SOACMapper frep trep m -> SOAC frep -> m (SOAC trep)
    mapSOACM SOACMapper rep rep Identity
    substitute
        where
          substitute :: SOACMapper rep rep Identity
    substitute =
            SOACMapper
              { mapOnSOACSubExp :: SubExp -> Identity SubExp
    mapOnSOACSubExp = SubExp -> Identity SubExp
    forall a. a -> Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (SubExp -> Identity SubExp)
    -> (SubExp -> SubExp) -> SubExp -> Identity SubExp
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Map VName VName -> SubExp -> SubExp
    forall a. Substitute a => Map VName VName -> a -> a
    substituteNames Map VName VName
    subst,
                mapOnSOACLambda :: Lambda rep -> Identity (Lambda rep)
    mapOnSOACLambda = Lambda rep -> Identity (Lambda rep)
    forall a. a -> Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Lambda rep -> Identity (Lambda rep))
    -> (Lambda rep -> Lambda rep)
    -> Lambda rep
    -> Identity (Lambda rep)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Map VName VName -> Lambda rep -> Lambda rep
    forall a. Substitute a => Map VName VName -> a -> a
    substituteNames Map VName VName
    subst,
                mapOnSOACVName :: VName -> Identity VName
    mapOnSOACVName = VName -> Identity VName
    forall a. a -> Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (VName -> Identity VName)
    -> (VName -> VName) -> VName -> Identity VName
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Map VName VName -> VName -> VName
    forall a. Substitute a => Map VName VName -> a -> a
    substituteNames Map VName VName
    subst
              }

    instance (ASTRep rep) => Rename (SOAC rep) where
      rename :: SOAC rep -> RenameM (SOAC rep)
    rename = SOACMapper rep rep RenameM -> SOAC rep -> RenameM (SOAC rep)
    forall (m :: * -> *) frep trep.
    Monad m =>
    SOACMapper frep trep m -> SOAC frep -> m (SOAC trep)
    mapSOACM SOACMapper rep rep RenameM
    renamer
        where
          renamer :: SOACMapper rep rep RenameM
    renamer = (SubExp -> RenameM SubExp)
    -> (Lambda rep -> RenameM (Lambda rep))
    -> (VName -> RenameM VName)
    -> SOACMapper rep rep RenameM
    forall frep trep (m :: * -> *).
    (SubExp -> m SubExp)
    -> (Lambda frep -> m (Lambda trep))
    -> (VName -> m VName)
    -> SOACMapper frep trep m
    SOACMapper SubExp -> RenameM SubExp
    forall a. Rename a => a -> RenameM a
    rename Lambda rep -> RenameM (Lambda rep)
    forall a. Rename a => a -> RenameM a
    rename VName -> RenameM VName
    forall a. Rename a => a -> RenameM a
    rename

    -- | The type of a SOAC.
    soacType :: (Typed (LParamInfo rep)) => SOAC rep -> [Type]
    soacType :: forall rep. Typed (LParamInfo rep) => SOAC rep -> [Type]
    soacType (JVP [SubExp]
    _ [SubExp]
    _ Lambda rep
    lam) =
      Lambda rep -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda rep
    lam [Type] -> [Type] -> [Type]
    forall a. [a] -> [a] -> [a]
    ++ Lambda rep -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda rep
    lam
    soacType (VJP [SubExp]
    _ [SubExp]
    _ Lambda rep
    lam) =
      Lambda rep -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda rep
    lam [Type] -> [Type] -> [Type]
    forall a. [a] -> [a] -> [a]
    ++ (Param (LParamInfo rep) -> Type)
    -> [Param (LParamInfo rep)] -> [Type]
    forall a b. (a -> b) -> [a] -> [b]
    map Param (LParamInfo rep) -> Type
    forall dec. Typed dec => Param dec -> Type
    paramType (Lambda rep -> [Param (LParamInfo rep)]
    forall {rep}. Lambda rep -> [Param (LParamInfo rep)]
    lambdaParams Lambda rep
    lam)
    soacType (Stream SubExp
    outersize [VName]
    _ [SubExp]
    accs Lambda rep
    lam) =
      (Type -> Type) -> [Type] -> [Type]
    forall a b. (a -> b) -> [a] -> [b]
    map (Map VName SubExp -> Type -> Type
    substNamesInType Map VName SubExp
    substs) [Type]
    rtp
      where
        nms :: [VName]
    nms = (Param (LParamInfo rep) -> VName)
    -> [Param (LParamInfo rep)] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map Param (LParamInfo rep) -> VName
    forall dec. Param dec -> VName
    paramName ([Param (LParamInfo rep)] -> [VName])
    -> [Param (LParamInfo rep)] -> [VName]
    forall a b. (a -> b) -> a -> b
    $ Int -> [Param (LParamInfo rep)] -> [Param (LParamInfo rep)]
    forall a. Int -> [a] -> [a]
    take (Int
    1 Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + [SubExp] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [SubExp]
    accs) [Param (LParamInfo rep)]
    params
        substs :: Map VName SubExp
    substs = [(VName, SubExp)] -> Map VName SubExp
    forall k a. Ord k => [(k, a)] -> Map k a
    M.fromList ([(VName, SubExp)] -> Map VName SubExp)
    -> [(VName, SubExp)] -> Map VName SubExp
    forall a b. (a -> b) -> a -> b
    $ [VName] -> [SubExp] -> [(VName, SubExp)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [VName]
    nms (SubExp
    outersize SubExp -> [SubExp] -> [SubExp]
    forall a. a -> [a] -> [a]
    : [SubExp]
    accs)
        Lambda [Param (LParamInfo rep)]
    params [Type]
    rtp Body rep
    _ = Lambda rep
    lam
    soacType (Scatter SubExp
    _w [VName]
    _ivs ScatterSpec VName
    dests Lambda rep
    lam) =
      (Type -> Shape -> Type) -> [Type] -> [Shape] -> [Type]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith Type -> Shape -> Type
    arrayOfShape (([([Type], Type)] -> Type) -> [[([Type], Type)]] -> [Type]
    forall a b. (a -> b) -> [a] -> [b]
    map (([Type], Type) -> Type
    forall a b. (a, b) -> b
    snd (([Type], Type) -> Type)
    -> ([([Type], Type)] -> ([Type], Type)) -> [([Type], Type)] -> Type
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . [([Type], Type)] -> ([Type], Type)
    forall a. HasCallStack => [a] -> a
    head) [[([Type], Type)]]
    rets) [Shape]
    shapes
      where
        ([Shape]
    shapes, [VName]
    _, [[([Type], Type)]]
    rets) =
          [(Shape, VName, [([Type], Type)])]
    -> ([Shape], [VName], [[([Type], Type)]])
    forall a b c. [(a, b, c)] -> ([a], [b], [c])
    unzip3 ([(Shape, VName, [([Type], Type)])]
     -> ([Shape], [VName], [[([Type], Type)]]))
    -> [(Shape, VName, [([Type], Type)])]
    -> ([Shape], [VName], [[([Type], Type)]])
    forall a b. (a -> b) -> a -> b
    $ ScatterSpec VName -> [Type] -> [(Shape, VName, [([Type], Type)])]
    forall array a.
    ScatterSpec array -> [a] -> [(Shape, array, [([a], a)])]
    groupScatterResults ScatterSpec VName
    dests ([Type] -> [(Shape, VName, [([Type], Type)])])
    -> [Type] -> [(Shape, VName, [([Type], Type)])]
    forall a b. (a -> b) -> a -> b
    $ Lambda rep -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda rep
    lam
    soacType (Hist SubExp
    _ [VName]
    _ [HistOp rep]
    ops Lambda rep
    _bucket_fun) = do
      HistOp rep
    op <- [HistOp rep]
    ops
      (Type -> Type) -> [Type] -> [Type]
    forall a b. (a -> b) -> [a] -> [b]
    map (Type -> Shape -> Type
    `arrayOfShape` HistOp rep -> Shape
    forall rep. HistOp rep -> Shape
    histShape HistOp rep
    op) (Lambda rep -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType (Lambda rep -> [Type]) -> Lambda rep -> [Type]
    forall a b. (a -> b) -> a -> b
    $ HistOp rep -> Lambda rep
    forall rep. HistOp rep -> Lambda rep
    histOp HistOp rep
    op)
    soacType (Screma SubExp
    w [VName]
    _arrs ScremaForm rep
    form) =
      SubExp -> ScremaForm rep -> [Type]
    forall rep. SubExp -> ScremaForm rep -> [Type]
    scremaType SubExp
    w ScremaForm rep
    form

    instance TypedOp SOAC where
      opType :: forall rep (m :: * -> *). HasScope rep m => SOAC rep -> m [ExtType]
    opType = [ExtType] -> m [ExtType]
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure ([ExtType] -> m [ExtType])
    -> (SOAC rep -> [ExtType]) -> SOAC rep -> m [ExtType]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . [Type] -> [ExtType]
    forall u. [TypeBase Shape u] -> [TypeBase ExtShape u]
    staticShapes ([Type] -> [ExtType])
    -> (SOAC rep -> [Type]) -> SOAC rep -> [ExtType]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . SOAC rep -> [Type]
    forall rep. Typed (LParamInfo rep) => SOAC rep -> [Type]
    soacType

    instance AliasedOp SOAC where
      opAliases :: forall rep. Aliased rep => SOAC rep -> [Names]
    opAliases = (Type -> Names) -> [Type] -> [Names]
    forall a b. (a -> b) -> [a] -> [b]
    map (Names -> Type -> Names
    forall a b. a -> b -> a
    const Names
    forall a. Monoid a => a
    mempty) ([Type] -> [Names]) -> (SOAC rep -> [Type]) -> SOAC rep -> [Names]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . SOAC rep -> [Type]
    forall rep. Typed (LParamInfo rep) => SOAC rep -> [Type]
    soacType

      consumedInOp :: forall rep. Aliased rep => SOAC rep -> Names
    consumedInOp JVP {} = Names
    forall a. Monoid a => a
    mempty
      consumedInOp VJP {} = Names
    forall a. Monoid a => a
    mempty
      -- Only map functions can consume anything.  The operands to scan
      -- and reduce functions are always considered "fresh".
      consumedInOp (Screma SubExp
    _ [VName]
    arrs (ScremaForm Lambda rep
    map_lam [Scan rep]
    _ [Reduce rep]
    _)) =
        (VName -> VName) -> Names -> Names
    mapNames VName -> VName
    consumedArray (Names -> Names) -> Names -> Names
    forall a b. (a -> b) -> a -> b
    $ Lambda rep -> Names
    forall rep. Aliased rep => Lambda rep -> Names
    consumedByLambda Lambda rep
    map_lam
        where
          consumedArray :: VName -> VName
    consumedArray VName
    v = VName -> Maybe VName -> VName
    forall a. a -> Maybe a -> a
    fromMaybe VName
    v (Maybe VName -> VName) -> Maybe VName -> VName
    forall a b. (a -> b) -> a -> b
    $ VName -> [(VName, VName)] -> Maybe VName
    forall a b. Eq a => a -> [(a, b)] -> Maybe b
    lookup VName
    v [(VName, VName)]
    params_to_arrs
          params_to_arrs :: [(VName, VName)]
    params_to_arrs = [VName] -> [VName] -> [(VName, VName)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip ((Param (LParamInfo rep) -> VName)
    -> [Param (LParamInfo rep)] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map Param (LParamInfo rep) -> VName
    forall dec. Param dec -> VName
    paramName ([Param (LParamInfo rep)] -> [VName])
    -> [Param (LParamInfo rep)] -> [VName]
    forall a b. (a -> b) -> a -> b
    $ Lambda rep -> [Param (LParamInfo rep)]
    forall {rep}. Lambda rep -> [Param (LParamInfo rep)]
    lambdaParams Lambda rep
    map_lam) [VName]
    arrs
      consumedInOp (Stream SubExp
    _ [VName]
    arrs [SubExp]
    accs Lambda rep
    lam) =
        [VName] -> Names
    namesFromList ([VName] -> Names) -> [VName] -> Names
    forall a b. (a -> b) -> a -> b
    $ [SubExp] -> [VName]
    subExpVars ([SubExp] -> [VName]) -> [SubExp] -> [VName]
    forall a b. (a -> b) -> a -> b
    $ (VName -> SubExp) -> [VName] -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map VName -> SubExp
    consumedArray ([VName] -> [SubExp]) -> [VName] -> [SubExp]
    forall a b. (a -> b) -> a -> b
    $ Names -> [VName]
    namesToList (Names -> [VName]) -> Names -> [VName]
    forall a b. (a -> b) -> a -> b
    $ Lambda rep -> Names
    forall rep. Aliased rep => Lambda rep -> Names
    consumedByLambda Lambda rep
    lam
        where
          consumedArray :: VName -> SubExp
    consumedArray VName
    v = SubExp -> Maybe SubExp -> SubExp
    forall a. a -> Maybe a -> a
    fromMaybe (VName -> SubExp
    Var VName
    v) (Maybe SubExp -> SubExp) -> Maybe SubExp -> SubExp
    forall a b. (a -> b) -> a -> b
    $ VName -> [(VName, SubExp)] -> Maybe SubExp
    forall a b. Eq a => a -> [(a, b)] -> Maybe b
    lookup VName
    v [(VName, SubExp)]
    paramsToInput
          -- Drop the chunk parameter, which cannot alias anything.
          paramsToInput :: [(VName, SubExp)]
    paramsToInput =
            [VName] -> [SubExp] -> [(VName, SubExp)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip ((Param (LParamInfo rep) -> VName)
    -> [Param (LParamInfo rep)] -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map Param (LParamInfo rep) -> VName
    forall dec. Param dec -> VName
    paramName ([Param (LParamInfo rep)] -> [VName])
    -> [Param (LParamInfo rep)] -> [VName]
    forall a b. (a -> b) -> a -> b
    $ Int -> [Param (LParamInfo rep)] -> [Param (LParamInfo rep)]
    forall a. Int -> [a] -> [a]
    drop Int
    1 ([Param (LParamInfo rep)] -> [Param (LParamInfo rep)])
    -> [Param (LParamInfo rep)] -> [Param (LParamInfo rep)]
    forall a b. (a -> b) -> a -> b
    $ Lambda rep -> [Param (LParamInfo rep)]
    forall {rep}. Lambda rep -> [Param (LParamInfo rep)]
    lambdaParams Lambda rep
    lam) ([SubExp]
    accs [SubExp] -> [SubExp] -> [SubExp]
    forall a. [a] -> [a] -> [a]
    ++ (VName -> SubExp) -> [VName] -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map VName -> SubExp
    Var [VName]
    arrs)
      consumedInOp (Scatter SubExp
    _ [VName]
    _ ScatterSpec VName
    spec Lambda rep
    _) =
        [VName] -> Names
    namesFromList ([VName] -> Names) -> [VName] -> Names
    forall a b. (a -> b) -> a -> b
    $ ((Shape, Int, VName) -> VName) -> ScatterSpec VName -> [VName]
    forall a b. (a -> b) -> [a] -> [b]
    map (\(Shape
    _, Int
    _, VName
    a) -> VName
    a) ScatterSpec VName
    spec
      consumedInOp (Hist SubExp
    _ [VName]
    _ [HistOp rep]
    ops Lambda rep
    _) =
        [VName] -> Names
    namesFromList ([VName] -> Names) -> [VName] -> Names
    forall a b. (a -> b) -> a -> b
    $ (HistOp rep -> [VName]) -> [HistOp rep] -> [VName]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap HistOp rep -> [VName]
    forall rep. HistOp rep -> [VName]
    histDest [HistOp rep]
    ops

    mapHistOp ::
      (Lambda frep -> Lambda trep) ->
      HistOp frep ->
      HistOp trep
    mapHistOp :: forall frep trep.
    (Lambda frep -> Lambda trep) -> HistOp frep -> HistOp trep
    mapHistOp Lambda frep -> Lambda trep
    f (HistOp Shape
    w SubExp
    rf [VName]
    dests [SubExp]
    nes Lambda frep
    lam) =
      Shape
    -> SubExp -> [VName] -> [SubExp] -> Lambda trep -> HistOp trep
    forall rep.
    Shape -> SubExp -> [VName] -> [SubExp] -> Lambda rep -> HistOp rep
    HistOp Shape
    w SubExp
    rf [VName]
    dests [SubExp]
    nes (Lambda trep -> HistOp trep) -> Lambda trep -> HistOp trep
    forall a b. (a -> b) -> a -> b
    $ Lambda frep -> Lambda trep
    f Lambda frep
    lam

    instance CanBeAliased SOAC where
      addOpAliases :: forall rep.
    AliasableRep rep =>
    AliasTable -> SOAC rep -> SOAC (Aliases rep)
    addOpAliases AliasTable
    aliases (JVP [SubExp]
    args [SubExp]
    vec Lambda rep
    lam) =
        [SubExp] -> [SubExp] -> Lambda (Aliases rep) -> SOAC (Aliases rep)
    forall rep. [SubExp] -> [SubExp] -> Lambda rep -> SOAC rep
    JVP [SubExp]
    args [SubExp]
    vec (AliasTable -> Lambda rep -> Lambda (Aliases rep)
    forall rep.
    AliasableRep rep =>
    AliasTable -> Lambda rep -> Lambda (Aliases rep)
    Alias.analyseLambda AliasTable
    aliases Lambda rep
    lam)
      addOpAliases AliasTable
    aliases (VJP [SubExp]
    args [SubExp]
    vec Lambda rep
    lam) =
        [SubExp] -> [SubExp] -> Lambda (Aliases rep) -> SOAC (Aliases rep)
    forall rep. [SubExp] -> [SubExp] -> Lambda rep -> SOAC rep
    VJP [SubExp]
    args [SubExp]
    vec (AliasTable -> Lambda rep -> Lambda (Aliases rep)
    forall rep.
    AliasableRep rep =>
    AliasTable -> Lambda rep -> Lambda (Aliases rep)
    Alias.analyseLambda AliasTable
    aliases Lambda rep
    lam)
      addOpAliases AliasTable
    aliases (Stream SubExp
    size [VName]
    arr [SubExp]
    accs Lambda rep
    lam) =
        SubExp
    -> [VName]
    -> [SubExp]
    -> Lambda (Aliases rep)
    -> SOAC (Aliases rep)
    forall rep. SubExp -> [VName] -> [SubExp] -> Lambda rep -> SOAC rep
    Stream SubExp
    size [VName]
    arr [SubExp]
    accs (Lambda (Aliases rep) -> SOAC (Aliases rep))
    -> Lambda (Aliases rep) -> SOAC (Aliases rep)
    forall a b. (a -> b) -> a -> b
    $ AliasTable -> Lambda rep -> Lambda (Aliases rep)
    forall rep.
    AliasableRep rep =>
    AliasTable -> Lambda rep -> Lambda (Aliases rep)
    Alias.analyseLambda AliasTable
    aliases Lambda rep
    lam
      addOpAliases AliasTable
    aliases (Scatter SubExp
    len [VName]
    arrs ScatterSpec VName
    dests Lambda rep
    lam) =
        SubExp
    -> [VName]
    -> ScatterSpec VName
    -> Lambda (Aliases rep)
    -> SOAC (Aliases rep)
    forall rep.
    SubExp -> [VName] -> ScatterSpec VName -> Lambda rep -> SOAC rep
    Scatter SubExp
    len [VName]
    arrs ScatterSpec VName
    dests (AliasTable -> Lambda rep -> Lambda (Aliases rep)
    forall rep.
    AliasableRep rep =>
    AliasTable -> Lambda rep -> Lambda (Aliases rep)
    Alias.analyseLambda AliasTable
    aliases Lambda rep
    lam)
      addOpAliases AliasTable
    aliases (Hist SubExp
    w [VName]
    arrs [HistOp rep]
    ops Lambda rep
    bucket_fun) =
        SubExp
    -> [VName]
    -> [HistOp (Aliases rep)]
    -> Lambda (Aliases rep)
    -> SOAC (Aliases rep)
    forall rep.
    SubExp -> [VName] -> [HistOp rep] -> Lambda rep -> SOAC rep
    Hist
          SubExp
    w
          [VName]
    arrs
          ((HistOp rep -> HistOp (Aliases rep))
    -> [HistOp rep] -> [HistOp (Aliases rep)]
    forall a b. (a -> b) -> [a] -> [b]
    map ((Lambda rep -> Lambda (Aliases rep))
    -> HistOp rep -> HistOp (Aliases rep)
    forall frep trep.
    (Lambda frep -> Lambda trep) -> HistOp frep -> HistOp trep
    mapHistOp (AliasTable -> Lambda rep -> Lambda (Aliases rep)
    forall rep.
    AliasableRep rep =>
    AliasTable -> Lambda rep -> Lambda (Aliases rep)
    Alias.analyseLambda AliasTable
    aliases)) [HistOp rep]
    ops)
          (AliasTable -> Lambda rep -> Lambda (Aliases rep)
    forall rep.
    AliasableRep rep =>
    AliasTable -> Lambda rep -> Lambda (Aliases rep)
    Alias.analyseLambda AliasTable
    aliases Lambda rep
    bucket_fun)
      addOpAliases AliasTable
    aliases (Screma SubExp
    w [VName]
    arrs (ScremaForm Lambda rep
    map_lam [Scan rep]
    scans [Reduce rep]
    reds)) =
        SubExp -> [VName] -> ScremaForm (Aliases rep) -> SOAC (Aliases rep)
    forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
    Screma SubExp
    w [VName]
    arrs (ScremaForm (Aliases rep) -> SOAC (Aliases rep))
    -> ScremaForm (Aliases rep) -> SOAC (Aliases rep)
    forall a b. (a -> b) -> a -> b
    $
          Lambda (Aliases rep)
    -> [Scan (Aliases rep)]
    -> [Reduce (Aliases rep)]
    -> ScremaForm (Aliases rep)
    forall rep.
    Lambda rep -> [Scan rep] -> [Reduce rep] -> ScremaForm rep
    ScremaForm
            (AliasTable -> Lambda rep -> Lambda (Aliases rep)
    forall rep.
    AliasableRep rep =>
    AliasTable -> Lambda rep -> Lambda (Aliases rep)
    Alias.analyseLambda AliasTable
    aliases Lambda rep
    map_lam)
            ((Scan rep -> Scan (Aliases rep))
    -> [Scan rep] -> [Scan (Aliases rep)]
    forall a b. (a -> b) -> [a] -> [b]
    map Scan rep -> Scan (Aliases rep)
    onScan [Scan rep]
    scans)
            ((Reduce rep -> Reduce (Aliases rep))
    -> [Reduce rep] -> [Reduce (Aliases rep)]
    forall a b. (a -> b) -> [a] -> [b]
    map Reduce rep -> Reduce (Aliases rep)
    onRed [Reduce rep]
    reds)
        where
          onRed :: Reduce rep -> Reduce (Aliases rep)
    onRed Reduce rep
    red = Reduce rep
    red {redLambda = Alias.analyseLambda aliases $ redLambda red}
          onScan :: Scan rep -> Scan (Aliases rep)
    onScan Scan rep
    scan = Scan rep
    scan {scanLambda = Alias.analyseLambda aliases $ scanLambda scan}

    instance IsOp SOAC where
      safeOp :: forall rep. ASTRep rep => SOAC rep -> Bool
    safeOp SOAC rep
    _ = Bool
    False
      cheapOp :: forall rep. ASTRep rep => SOAC rep -> Bool
    cheapOp SOAC rep
    _ = Bool
    False
      opDependencies :: forall rep. ASTRep rep => SOAC rep -> [Names]
    opDependencies (Stream SubExp
    w [VName]
    arrs [SubExp]
    accs Lambda rep
    lam) =
        let accs_deps :: [Names]
    accs_deps = (SubExp -> Names) -> [SubExp] -> [Names]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExp -> Names
    depsOf' [SubExp]
    accs
            arrs_deps :: [Names]
    arrs_deps = SubExp -> [VName] -> [Names]
    depsOfArrays SubExp
    w [VName]
    arrs
         in AliasTable -> Lambda rep -> [Names] -> [Names]
    forall rep.
    ASTRep rep =>
    AliasTable -> Lambda rep -> [Names] -> [Names]
    lambdaDependencies AliasTable
    forall a. Monoid a => a
    mempty Lambda rep
    lam ([Names]
    arrs_deps [Names] -> [Names] -> [Names]
    forall a. Semigroup a => a -> a -> a
    <> [Names]
    accs_deps)
      opDependencies (Hist SubExp
    w [VName]
    arrs [HistOp rep]
    ops Lambda rep
    lam) =
        let bucket_fun_deps' :: [Names]
    bucket_fun_deps' = AliasTable -> Lambda rep -> [Names] -> [Names]
    forall rep.
    ASTRep rep =>
    AliasTable -> Lambda rep -> [Names] -> [Names]
    lambdaDependencies AliasTable
    forall a. Monoid a => a
    mempty Lambda rep
    lam (SubExp -> [VName] -> [Names]
    depsOfArrays SubExp
    w [VName]
    arrs)
            -- Bucket function results are indices followed by values.
            -- Reshape this to align with list of histogram operations.
            ranks :: [Int]
    ranks = (HistOp rep -> Int) -> [HistOp rep] -> [Int]
    forall a b. (a -> b) -> [a] -> [b]
    map (Shape -> Int
    forall a. ArrayShape a => a -> Int
    shapeRank (Shape -> Int) -> (HistOp rep -> Shape) -> HistOp rep -> Int
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . HistOp rep -> Shape
    forall rep. HistOp rep -> Shape
    histShape) [HistOp rep]
    ops
            value_lengths :: [Int]
    value_lengths = (HistOp rep -> Int) -> [HistOp rep] -> [Int]
    forall a b. (a -> b) -> [a] -> [b]
    map ([SubExp] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length ([SubExp] -> Int) -> (HistOp rep -> [SubExp]) -> HistOp rep -> Int
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . HistOp rep -> [SubExp]
    forall rep. HistOp rep -> [SubExp]
    histNeutral) [HistOp rep]
    ops
            ([Names]
    indices, [Names]
    values) = Int -> [Names] -> ([Names], [Names])
    forall a. Int -> [a] -> ([a], [a])
    splitAt ([Int] -> Int
    forall a. Num a => [a] -> a
    forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
    sum [Int]
    ranks) [Names]
    bucket_fun_deps'
            bucket_fun_deps :: [[Names]]
    bucket_fun_deps =
              ([Names] -> [Names] -> [Names])
    -> [[Names]] -> [[Names]] -> [[Names]]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith
                [Names] -> [Names] -> [Names]
    forall {b}. Monoid b => [b] -> [b] -> [b]
    concatIndicesToEachValue
                ([Int] -> [Names] -> [[Names]]
    forall a. [Int] -> [a] -> [[a]]
    chunks [Int]
    ranks [Names]
    indices)
                ([Int] -> [Names] -> [[Names]]
    forall a. [Int] -> [a] -> [[a]]
    chunks [Int]
    value_lengths [Names]
    values)
         in [[Names]] -> [Names]
    forall a. Monoid a => [a] -> a
    mconcat ([[Names]] -> [Names]) -> [[Names]] -> [Names]
    forall a b. (a -> b) -> a -> b
    $ ([Names] -> [Names] -> [Names])
    -> [[Names]] -> [[Names]] -> [[Names]]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith ((Names -> Names -> Names) -> [Names] -> [Names] -> [Names]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    (<>)) [[Names]]
    bucket_fun_deps ((HistOp rep -> [Names]) -> [HistOp rep] -> [[Names]]
    forall a b. (a -> b) -> [a] -> [b]
    map HistOp rep -> [Names]
    forall {rep}. ASTRep rep => HistOp rep -> [Names]
    depsOfHistOp [HistOp rep]
    ops)
        where
          depsOfHistOp :: HistOp rep -> [Names]
    depsOfHistOp (HistOp Shape
    dest_shape SubExp
    rf [VName]
    dests [SubExp]
    nes Lambda rep
    op) =
            let shape_deps :: Names
    shape_deps = Shape -> Names
    depsOfShape Shape
    dest_shape
                in_deps :: [Names]
    in_deps = (VName -> Names) -> [VName] -> [Names]
    forall a b. (a -> b) -> [a] -> [b]
    map (\VName
    vn -> VName -> Names
    oneName VName
    vn Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> Names
    shape_deps Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> SubExp -> Names
    depsOf' SubExp
    rf) [VName]
    dests
             in AliasTable -> Lambda rep -> [SubExp] -> [Names] -> [Names]
    forall rep.
    ASTRep rep =>
    AliasTable -> Lambda rep -> [SubExp] -> [Names] -> [Names]
    reductionDependencies AliasTable
    forall a. Monoid a => a
    mempty Lambda rep
    op [SubExp]
    nes [Names]
    in_deps
          -- A histogram operation may use the same index for multiple values.
          concatIndicesToEachValue :: [b] -> [b] -> [b]
    concatIndicesToEachValue [b]
    is [b]
    vs =
            let is_flat :: b
    is_flat = [b] -> b
    forall a. Monoid a => [a] -> a
    mconcat [b]
    is
             in (b -> b) -> [b] -> [b]
    forall a b. (a -> b) -> [a] -> [b]
    map (b
    is_flat <>) [b]
    vs
      opDependencies (Scatter SubExp
    w [VName]
    arrs ScatterSpec VName
    outputs Lambda rep
    lam) =
        let deps :: [Names]
    deps = AliasTable -> Lambda rep -> [Names] -> [Names]
    forall rep.
    ASTRep rep =>
    AliasTable -> Lambda rep -> [Names] -> [Names]
    lambdaDependencies AliasTable
    forall a. Monoid a => a
    mempty Lambda rep
    lam (SubExp -> [VName] -> [Names]
    depsOfArrays SubExp
    w [VName]
    arrs)
         in ((Shape, VName, [([Names], Names)]) -> Names)
    -> [(Shape, VName, [([Names], Names)])] -> [Names]
    forall a b. (a -> b) -> [a] -> [b]
    map (Shape, VName, [([Names], Names)]) -> Names
    forall {a}. (a, VName, [([Names], Names)]) -> Names
    flattenBlocks (ScatterSpec VName
    -> [Names] -> [(Shape, VName, [([Names], Names)])]
    forall array a.
    ScatterSpec array -> [a] -> [(Shape, array, [([a], a)])]
    groupScatterResults ScatterSpec VName
    outputs [Names]
    deps)
        where
          flattenBlocks :: (a, VName, [([Names], Names)]) -> Names
    flattenBlocks (a
    _, VName
    arr, [([Names], Names)]
    ivs) =
            VName -> Names
    oneName VName
    arr Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> [Names] -> Names
    forall a. Monoid a => [a] -> a
    mconcat ((([Names], Names) -> Names) -> [([Names], Names)] -> [Names]
    forall a b. (a -> b) -> [a] -> [b]
    map ([Names] -> Names
    forall a. Monoid a => [a] -> a
    mconcat ([Names] -> Names)
    -> (([Names], Names) -> [Names]) -> ([Names], Names) -> Names
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . ([Names], Names) -> [Names]
    forall a b. (a, b) -> a
    fst) [([Names], Names)]
    ivs) Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> [Names] -> Names
    forall a. Monoid a => [a] -> a
    mconcat ((([Names], Names) -> Names) -> [([Names], Names)] -> [Names]
    forall a b. (a -> b) -> [a] -> [b]
    map ([Names], Names) -> Names
    forall a b. (a, b) -> b
    snd [([Names], Names)]
    ivs)
      opDependencies (JVP [SubExp]
    args [SubExp]
    vec Lambda rep
    lam) =
        [[Names]] -> [Names]
    forall a. Monoid a => [a] -> a
    mconcat ([[Names]] -> [Names]) -> [[Names]] -> [Names]
    forall a b. (a -> b) -> a -> b
    $
          Int -> [Names] -> [[Names]]
    forall a. Int -> a -> [a]
    replicate Int
    2 ([Names] -> [[Names]]) -> [Names] -> [[Names]]
    forall a b. (a -> b) -> a -> b
    $
            AliasTable -> Lambda rep -> [Names] -> [Names]
    forall rep.
    ASTRep rep =>
    AliasTable -> Lambda rep -> [Names] -> [Names]
    lambdaDependencies AliasTable
    forall a. Monoid a => a
    mempty Lambda rep
    lam ([Names] -> [Names]) -> [Names] -> [Names]
    forall a b. (a -> b) -> a -> b
    $
              (Names -> Names -> Names) -> [Names] -> [Names] -> [Names]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    (<>) ((SubExp -> Names) -> [SubExp] -> [Names]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExp -> Names
    depsOf' [SubExp]
    args) ((SubExp -> Names) -> [SubExp] -> [Names]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExp -> Names
    depsOf' [SubExp]
    vec)
      opDependencies (VJP [SubExp]
    args [SubExp]
    vec Lambda rep
    lam) =
        AliasTable -> Lambda rep -> [Names] -> [Names]
    forall rep.
    ASTRep rep =>
    AliasTable -> Lambda rep -> [Names] -> [Names]
    lambdaDependencies
          AliasTable
    forall a. Monoid a => a
    mempty
          Lambda rep
    lam
          ((Names -> Names -> Names) -> [Names] -> [Names] -> [Names]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    (<>) ((SubExp -> Names) -> [SubExp] -> [Names]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExp -> Names
    depsOf' [SubExp]
    args) ((SubExp -> Names) -> [SubExp] -> [Names]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExp -> Names
    depsOf' [SubExp]
    vec))
          [Names] -> [Names] -> [Names]
    forall a. Semigroup a => a -> a -> a
    <> (Param (LParamInfo rep) -> Names)
    -> [Param (LParamInfo rep)] -> [Names]
    forall a b. (a -> b) -> [a] -> [b]
    map (Names -> Param (LParamInfo rep) -> Names
    forall a b. a -> b -> a
    const (Names -> Param (LParamInfo rep) -> Names)
    -> Names -> Param (LParamInfo rep) -> Names
    forall a b. (a -> b) -> a -> b
    $ [SubExp] -> Names
    forall a. FreeIn a => a -> Names
    freeIn [SubExp]
    args Names -> Names -> Names
    forall a. Semigroup a => a -> a -> a
    <> Lambda rep -> Names
    forall a. FreeIn a => a -> Names
    freeIn Lambda rep
    lam) (Lambda rep -> [Param (LParamInfo rep)]
    forall {rep}. Lambda rep -> [Param (LParamInfo rep)]
    lambdaParams Lambda rep
    lam)
      opDependencies (Screma SubExp
    w [VName]
    arrs (ScremaForm Lambda rep
    map_lam [Scan rep]
    scans [Reduce rep]
    reds)) =
        let ([Names]
    scans_in, [Names]
    reds_in, [Names]
    map_deps) =
              Int -> Int -> [Names] -> ([Names], [Names], [Names])
    forall a. Int -> Int -> [a] -> ([a], [a], [a])
    splitAt3 ([Scan rep] -> Int
    forall rep. [Scan rep] -> Int
    scanResults [Scan rep]
    scans) ([Reduce rep] -> Int
    forall rep. [Reduce rep] -> Int
    redResults [Reduce rep]
    reds) ([Names] -> ([Names], [Names], [Names]))
    -> [Names] -> ([Names], [Names], [Names])
    forall a b. (a -> b) -> a -> b
    $
                AliasTable -> Lambda rep -> [Names] -> [Names]
    forall rep.
    ASTRep rep =>
    AliasTable -> Lambda rep -> [Names] -> [Names]
    lambdaDependencies AliasTable
    forall a. Monoid a => a
    mempty Lambda rep
    map_lam (SubExp -> [VName] -> [Names]
    depsOfArrays SubExp
    w [VName]
    arrs)
            scans_deps :: [Names]
    scans_deps =
              ((Scan rep, [Names]) -> [Names])
    -> [(Scan rep, [Names])] -> [Names]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap (Scan rep, [Names]) -> [Names]
    forall {rep}. ASTRep rep => (Scan rep, [Names]) -> [Names]
    depsOfScan ([Scan rep] -> [[Names]] -> [(Scan rep, [Names])]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Scan rep]
    scans ([[Names]] -> [(Scan rep, [Names])])
    -> [[Names]] -> [(Scan rep, [Names])]
    forall a b. (a -> b) -> a -> b
    $ [Int] -> [Names] -> [[Names]]
    forall a. [Int] -> [a] -> [[a]]
    chunks ([Scan rep] -> [Int]
    forall rep. [Scan rep] -> [Int]
    scanSizes [Scan rep]
    scans) [Names]
    scans_in)
            reds_deps :: [Names]
    reds_deps =
              ((Reduce rep, [Names]) -> [Names])
    -> [(Reduce rep, [Names])] -> [Names]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap (Reduce rep, [Names]) -> [Names]
    forall {rep}. ASTRep rep => (Reduce rep, [Names]) -> [Names]
    depsOfRed ([Reduce rep] -> [[Names]] -> [(Reduce rep, [Names])]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Reduce rep]
    reds ([[Names]] -> [(Reduce rep, [Names])])
    -> [[Names]] -> [(Reduce rep, [Names])]
    forall a b. (a -> b) -> a -> b
    $ [Int] -> [Names] -> [[Names]]
    forall a. [Int] -> [a] -> [[a]]
    chunks ([Reduce rep] -> [Int]
    forall rep. [Reduce rep] -> [Int]
    redSizes [Reduce rep]
    reds) [Names]
    reds_in)
         in [Names]
    scans_deps [Names] -> [Names] -> [Names]
    forall a. Semigroup a => a -> a -> a
    <> [Names]
    reds_deps [Names] -> [Names] -> [Names]
    forall a. Semigroup a => a -> a -> a
    <> [Names]
    map_deps
        where
          depsOfScan :: (Scan rep, [Names]) -> [Names]
    depsOfScan (Scan Lambda rep
    lam [SubExp]
    nes, [Names]
    deps_in) =
            AliasTable -> Lambda rep -> [SubExp] -> [Names] -> [Names]
    forall rep.
    ASTRep rep =>
    AliasTable -> Lambda rep -> [SubExp] -> [Names] -> [Names]
    reductionDependencies AliasTable
    forall a. Monoid a => a
    mempty Lambda rep
    lam [SubExp]
    nes [Names]
    deps_in
          depsOfRed :: (Reduce rep, [Names]) -> [Names]
    depsOfRed (Reduce Commutativity
    _ Lambda rep
    lam [SubExp]
    nes, [Names]
    deps_in) =
            AliasTable -> Lambda rep -> [SubExp] -> [Names] -> [Names]
    forall rep.
    ASTRep rep =>
    AliasTable -> Lambda rep -> [SubExp] -> [Names] -> [Names]
    reductionDependencies AliasTable
    forall a. Monoid a => a
    mempty Lambda rep
    lam [SubExp]
    nes [Names]
    deps_in

    substNamesInType :: M.Map VName SubExp -> Type -> Type
    substNamesInType :: Map VName SubExp -> Type -> Type
    substNamesInType Map VName SubExp
    _ t :: Type
    t@Prim {} = Type
    t
    substNamesInType Map VName SubExp
    _ t :: Type
    t@Acc {} = Type
    t
    substNamesInType Map VName SubExp
    _ (Mem Space
    space) = Space -> Type
    forall shape u. Space -> TypeBase shape u
    Mem Space
    space
    substNamesInType Map VName SubExp
    subs (Array PrimType
    btp Shape
    shp NoUniqueness
    u) =
      let shp' :: Shape
    shp' = [SubExp] -> Shape
    forall d. [d] -> ShapeBase d
    Shape ([SubExp] -> Shape) -> [SubExp] -> Shape
    forall a b. (a -> b) -> a -> b
    $ (SubExp -> SubExp) -> [SubExp] -> [SubExp]
    forall a b. (a -> b) -> [a] -> [b]
    map (Map VName SubExp -> SubExp -> SubExp
    substNamesInSubExp Map VName SubExp
    subs) (Shape -> [SubExp]
    forall d. ShapeBase d -> [d]
    shapeDims Shape
    shp)
       in PrimType -> Shape -> NoUniqueness -> Type
    forall shape u. PrimType -> shape -> u -> TypeBase shape u
    Array PrimType
    btp Shape
    shp' NoUniqueness
    u

    substNamesInSubExp :: M.Map VName SubExp -> SubExp -> SubExp
    substNamesInSubExp :: Map VName SubExp -> SubExp -> SubExp
    substNamesInSubExp Map VName SubExp
    _ e :: SubExp
    e@(Constant PrimValue
    _) = SubExp
    e
    substNamesInSubExp Map VName SubExp
    subs (Var VName
    idd) =
      SubExp -> VName -> Map VName SubExp -> SubExp
    forall k a. Ord k => a -> k -> Map k a -> a
    M.findWithDefault (VName -> SubExp
    Var VName
    idd) VName
    idd Map VName SubExp
    subs

    instance CanBeWise SOAC where
      addOpWisdom :: forall rep. Informing rep => SOAC rep -> SOAC (Wise rep)
    addOpWisdom = Identity (SOAC (Wise rep)) -> SOAC (Wise rep)
    forall a. Identity a -> a
    runIdentity (Identity (SOAC (Wise rep)) -> SOAC (Wise rep))
    -> (SOAC rep -> Identity (SOAC (Wise rep)))
    -> SOAC rep
    -> SOAC (Wise rep)
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . SOACMapper rep (Wise rep) Identity
    -> SOAC rep -> Identity (SOAC (Wise rep))
    forall (m :: * -> *) frep trep.
    Monad m =>
    SOACMapper frep trep m -> SOAC frep -> m (SOAC trep)
    mapSOACM ((SubExp -> Identity SubExp)
    -> (Lambda rep -> Identity (Lambda (Wise rep)))
    -> (VName -> Identity VName)
    -> SOACMapper rep (Wise rep) Identity
    forall frep trep (m :: * -> *).
    (SubExp -> m SubExp)
    -> (Lambda frep -> m (Lambda trep))
    -> (VName -> m VName)
    -> SOACMapper frep trep m
    SOACMapper SubExp -> Identity SubExp
    forall a. a -> Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Lambda (Wise rep) -> Identity (Lambda (Wise rep))
    forall a. a -> Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Lambda (Wise rep) -> Identity (Lambda (Wise rep)))
    -> (Lambda rep -> Lambda (Wise rep))
    -> Lambda rep
    -> Identity (Lambda (Wise rep))
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Lambda rep -> Lambda (Wise rep)
    forall rep. Informing rep => Lambda rep -> Lambda (Wise rep)
    informLambda) VName -> Identity VName
    forall a. a -> Identity a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure)

    instance (RepTypes rep) => ST.IndexOp (SOAC rep) where
      indexOp :: forall rep.
    (ASTRep rep, IndexOp (Op rep)) =>
    SymbolTable rep
    -> Int -> SOAC rep -> [TPrimExp Int64 VName] -> Maybe Indexed
    indexOp SymbolTable rep
    vtable Int
    k SOAC rep
    soac [TPrimExp Int64 VName
    i] = do
        (Lambda rep
    lam, SubExpRes
    se, [Param (LParamInfo rep)]
    arr_params, [VName]
    arrs) <- SOAC rep
    -> Maybe (Lambda rep, SubExpRes, [Param (LParamInfo rep)], [VName])
    lambdaAndSubExp SOAC rep
    soac
        let arr_indexes :: Map VName (PrimExp VName, Certs)
    arr_indexes = [(VName, (PrimExp VName, Certs))]
    -> Map VName (PrimExp VName, Certs)
    forall k a. Ord k => [(k, a)] -> Map k a
    M.fromList ([(VName, (PrimExp VName, Certs))]
     -> Map VName (PrimExp VName, Certs))
    -> [(VName, (PrimExp VName, Certs))]
    -> Map VName (PrimExp VName, Certs)
    forall a b. (a -> b) -> a -> b
    $ [Maybe (VName, (PrimExp VName, Certs))]
    -> [(VName, (PrimExp VName, Certs))]
    forall a. [Maybe a] -> [a]
    catMaybes ([Maybe (VName, (PrimExp VName, Certs))]
     -> [(VName, (PrimExp VName, Certs))])
    -> [Maybe (VName, (PrimExp VName, Certs))]
    -> [(VName, (PrimExp VName, Certs))]
    forall a b. (a -> b) -> a -> b
    $ (Param (LParamInfo rep)
     -> VName -> Maybe (VName, (PrimExp VName, Certs)))
    -> [Param (LParamInfo rep)]
    -> [VName]
    -> [Maybe (VName, (PrimExp VName, Certs))]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith Param (LParamInfo rep)
    -> VName -> Maybe (VName, (PrimExp VName, Certs))
    arrIndex [Param (LParamInfo rep)]
    arr_params [VName]
    arrs
            arr_indexes' :: Map VName (PrimExp VName, Certs)
    arr_indexes' = (Map VName (PrimExp VName, Certs)
     -> Stm rep -> Map VName (PrimExp VName, Certs))
    -> Map VName (PrimExp VName, Certs)
    -> Seq (Stm rep)
    -> Map VName (PrimExp VName, Certs)
    forall b a. (b -> a -> b) -> b -> Seq a -> b
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl Map VName (PrimExp VName, Certs)
    -> Stm rep -> Map VName (PrimExp VName, Certs)
    expandPrimExpTable Map VName (PrimExp VName, Certs)
    arr_indexes (Seq (Stm rep) -> Map VName (PrimExp VName, Certs))
    -> Seq (Stm rep) -> Map VName (PrimExp VName, Certs)
    forall a b. (a -> b) -> a -> b
    $ Body rep -> Seq (Stm rep)
    forall rep. Body rep -> Stms rep
    bodyStms (Body rep -> Seq (Stm rep)) -> Body rep -> Seq (Stm rep)
    forall a b. (a -> b) -> a -> b
    $ Lambda rep -> Body rep
    forall rep. Lambda rep -> Body rep
    lambdaBody Lambda rep
    lam
        case SubExpRes
    se of
          SubExpRes Certs
    _ (Var VName
    v) -> (PrimExp VName -> Certs -> Indexed)
    -> (PrimExp VName, Certs) -> Indexed
    forall a b c. (a -> b -> c) -> (a, b) -> c
    uncurry ((Certs -> PrimExp VName -> Indexed)
    -> PrimExp VName -> Certs -> Indexed
    forall a b c. (a -> b -> c) -> b -> a -> c
    flip Certs -> PrimExp VName -> Indexed
    ST.Indexed) ((PrimExp VName, Certs) -> Indexed)
    -> Maybe (PrimExp VName, Certs) -> Maybe Indexed
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> VName
    -> Map VName (PrimExp VName, Certs) -> Maybe (PrimExp VName, Certs)
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    v Map VName (PrimExp VName, Certs)
    arr_indexes'
          SubExpRes
    _ -> Maybe Indexed
    forall a. Maybe a
    Nothing
        where
          lambdaAndSubExp :: SOAC rep
    -> Maybe (Lambda rep, SubExpRes, [Param (LParamInfo rep)], [VName])
    lambdaAndSubExp (Screma SubExp
    _ [VName]
    arrs (ScremaForm Lambda rep
    map_lam [Scan rep]
    scans [Reduce rep]
    reds)) =
            Int
    -> Lambda rep
    -> [VName]
    -> Maybe (Lambda rep, SubExpRes, [Param (LParamInfo rep)], [VName])
    nthMapOut ([Scan rep] -> Int
    forall rep. [Scan rep] -> Int
    scanResults [Scan rep]
    scans Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + [Reduce rep] -> Int
    forall rep. [Reduce rep] -> Int
    redResults [Reduce rep]
    reds) Lambda rep
    map_lam [VName]
    arrs
          lambdaAndSubExp SOAC rep
    _ =
            Maybe (Lambda rep, SubExpRes, [Param (LParamInfo rep)], [VName])
    forall a. Maybe a
    Nothing

          nthMapOut :: Int
    -> Lambda rep
    -> [VName]
    -> Maybe (Lambda rep, SubExpRes, [Param (LParamInfo rep)], [VName])
    nthMapOut Int
    num_accs Lambda rep
    lam [VName]
    arrs = do
            SubExpRes
    se <- Int -> Result -> Maybe SubExpRes
    forall int a. Integral int => int -> [a] -> Maybe a
    maybeNth (Int
    num_accs Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + Int
    k) (Result -> Maybe SubExpRes) -> Result -> Maybe SubExpRes
    forall a b. (a -> b) -> a -> b
    $ Body rep -> Result
    forall rep. Body rep -> Result
    bodyResult (Body rep -> Result) -> Body rep -> Result
    forall a b. (a -> b) -> a -> b
    $ Lambda rep -> Body rep
    forall rep. Lambda rep -> Body rep
    lambdaBody Lambda rep
    lam
            (Lambda rep, SubExpRes, [Param (LParamInfo rep)], [VName])
    -> Maybe (Lambda rep, SubExpRes, [Param (LParamInfo rep)], [VName])
    forall a. a -> Maybe a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Lambda rep
    lam, SubExpRes
    se, Int -> [Param (LParamInfo rep)] -> [Param (LParamInfo rep)]
    forall a. Int -> [a] -> [a]
    drop Int
    num_accs ([Param (LParamInfo rep)] -> [Param (LParamInfo rep)])
    -> [Param (LParamInfo rep)] -> [Param (LParamInfo rep)]
    forall a b. (a -> b) -> a -> b
    $ Lambda rep -> [Param (LParamInfo rep)]
    forall {rep}. Lambda rep -> [Param (LParamInfo rep)]
    lambdaParams Lambda rep
    lam, [VName]
    arrs)

          arrIndex :: Param (LParamInfo rep)
    -> VName -> Maybe (VName, (PrimExp VName, Certs))
    arrIndex Param (LParamInfo rep)
    p VName
    arr = do
            ST.Indexed Certs
    cs PrimExp VName
    pe <- VName -> [TPrimExp Int64 VName] -> SymbolTable rep -> Maybe Indexed
    forall rep.
    VName -> [TPrimExp Int64 VName] -> SymbolTable rep -> Maybe Indexed
    ST.index' VName
    arr [TPrimExp Int64 VName
    i] SymbolTable rep
    vtable
            (VName, (PrimExp VName, Certs))
    -> Maybe (VName, (PrimExp VName, Certs))
    forall a. a -> Maybe a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (Param (LParamInfo rep) -> VName
    forall dec. Param dec -> VName
    paramName Param (LParamInfo rep)
    p, (PrimExp VName
    pe, Certs
    cs))

          expandPrimExpTable :: Map VName (PrimExp VName, Certs)
    -> Stm rep -> Map VName (PrimExp VName, Certs)
    expandPrimExpTable Map VName (PrimExp VName, Certs)
    table Stm rep
    stm
            | [VName
    v] <- Pat (LetDec rep) -> [VName]
    forall dec. Pat dec -> [VName]
    patNames (Pat (LetDec rep) -> [VName]) -> Pat (LetDec rep) -> [VName]
    forall a b. (a -> b) -> a -> b
    $ Stm rep -> Pat (LetDec rep)
    forall rep. Stm rep -> Pat (LetDec rep)
    stmPat Stm rep
    stm,
              Just (PrimExp VName
    pe, Certs
    cs) <-
                WriterT Certs Maybe (PrimExp VName) -> Maybe (PrimExp VName, Certs)
    forall w (m :: * -> *) a. WriterT w m a -> m (a, w)
    runWriterT (WriterT Certs Maybe (PrimExp VName)
     -> Maybe (PrimExp VName, Certs))
    -> WriterT Certs Maybe (PrimExp VName)
    -> Maybe (PrimExp VName, Certs)
    forall a b. (a -> b) -> a -> b
    $ (VName -> WriterT Certs Maybe (PrimExp VName))
    -> Exp rep -> WriterT Certs Maybe (PrimExp VName)
    forall (m :: * -> *) rep v.
    (MonadFail m, RepTypes rep) =>
    (VName -> m (PrimExp v)) -> Exp rep -> m (PrimExp v)
    primExpFromExp (Map VName (PrimExp VName, Certs)
    -> VName -> WriterT Certs Maybe (PrimExp VName)
    asPrimExp Map VName (PrimExp VName, Certs)
    table) (Exp rep -> WriterT Certs Maybe (PrimExp VName))
    -> Exp rep -> WriterT Certs Maybe (PrimExp VName)
    forall a b. (a -> b) -> a -> b
    $ Stm rep -> Exp rep
    forall rep. Stm rep -> Exp rep
    stmExp Stm rep
    stm,
              (VName -> Bool) -> [VName] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all (VName -> SymbolTable rep -> Bool
    forall rep. VName -> SymbolTable rep -> Bool
    `ST.elem` SymbolTable rep
    vtable) (Certs -> [VName]
    unCerts (Certs -> [VName]) -> Certs -> [VName]
    forall a b. (a -> b) -> a -> b
    $ Stm rep -> Certs
    forall rep. Stm rep -> Certs
    stmCerts Stm rep
    stm) =
                VName
    -> (PrimExp VName, Certs)
    -> Map VName (PrimExp VName, Certs)
    -> Map VName (PrimExp VName, Certs)
    forall k a. Ord k => k -> a -> Map k a -> Map k a
    M.insert VName
    v (PrimExp VName
    pe, Stm rep -> Certs
    forall rep. Stm rep -> Certs
    stmCerts Stm rep
    stm Certs -> Certs -> Certs
    forall a. Semigroup a => a -> a -> a
    <> Certs
    cs) Map VName (PrimExp VName, Certs)
    table
            | Bool
    otherwise =
                Map VName (PrimExp VName, Certs)
    table

          asPrimExp :: Map VName (PrimExp VName, Certs)
    -> VName -> WriterT Certs Maybe (PrimExp VName)
    asPrimExp Map VName (PrimExp VName, Certs)
    table VName
    v
            | Just (PrimExp VName
    e, Certs
    cs) <- VName
    -> Map VName (PrimExp VName, Certs) -> Maybe (PrimExp VName, Certs)
    forall k a. Ord k => k -> Map k a -> Maybe a
    M.lookup VName
    v Map VName (PrimExp VName, Certs)
    table = Certs -> WriterT Certs Maybe ()
    forall w (m :: * -> *). MonadWriter w m => w -> m ()
    tell Certs
    cs WriterT Certs Maybe ()
    -> WriterT Certs Maybe (PrimExp VName)
    -> WriterT Certs Maybe (PrimExp VName)
    forall a b.
    WriterT Certs Maybe a
    -> WriterT Certs Maybe b -> WriterT Certs Maybe b
    forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
    >> PrimExp VName -> WriterT Certs Maybe (PrimExp VName)
    forall a. a -> WriterT Certs Maybe a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure PrimExp VName
    e
            | Just (Prim PrimType
    pt) <- VName -> SymbolTable rep -> Maybe Type
    forall rep. ASTRep rep => VName -> SymbolTable rep -> Maybe Type
    ST.lookupType VName
    v SymbolTable rep
    vtable =
                PrimExp VName -> WriterT Certs Maybe (PrimExp VName)
    forall a. a -> WriterT Certs Maybe a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure (PrimExp VName -> WriterT Certs Maybe (PrimExp VName))
    -> PrimExp VName -> WriterT Certs Maybe (PrimExp VName)
    forall a b. (a -> b) -> a -> b
    $ VName -> PrimType -> PrimExp VName
    forall v. v -> PrimType -> PrimExp v
    LeafExp VName
    v PrimType
    pt
            | Bool
    otherwise = Maybe (PrimExp VName) -> WriterT Certs Maybe (PrimExp VName)
    forall (m :: * -> *) a. Monad m => m a -> WriterT Certs m a
    forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
    (MonadTrans t, Monad m) =>
    m a -> t m a
    lift Maybe (PrimExp VName)
    forall a. Maybe a
    Nothing
      indexOp SymbolTable rep
    _ Int
    _ SOAC rep
    _ [TPrimExp Int64 VName]
    _ = Maybe Indexed
    forall a. Maybe a
    Nothing

    -- | Type-check a SOAC.
    typeCheckSOAC :: (TC.Checkable rep) => SOAC (Aliases rep) -> TC.TypeM rep ()
    typeCheckSOAC :: forall rep. Checkable rep => SOAC (Aliases rep) -> TypeM rep ()
    typeCheckSOAC (VJP [SubExp]
    args [SubExp]
    vec Lambda (Aliases rep)
    lam) = do
      [Arg]
    args' <- (SubExp -> TypeM rep Arg) -> [SubExp] -> TypeM rep [Arg]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM SubExp -> TypeM rep Arg
    forall rep. Checkable rep => SubExp -> TypeM rep Arg
    TC.checkArg [SubExp]
    args
      Lambda (Aliases rep) -> [Arg] -> TypeM rep ()
    forall rep.
    Checkable rep =>
    Lambda (Aliases rep) -> [Arg] -> TypeM rep ()
    TC.checkLambda Lambda (Aliases rep)
    lam ([Arg] -> TypeM rep ()) -> [Arg] -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $ (Arg -> Arg) -> [Arg] -> [Arg]
    forall a b. (a -> b) -> [a] -> [b]
    map Arg -> Arg
    TC.noArgAliases [Arg]
    args'
      [Type]
    vec_ts <- (SubExp -> TypeM rep Type) -> [SubExp] -> TypeM rep [Type]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM SubExp -> TypeM rep Type
    forall rep. Checkable rep => SubExp -> TypeM rep Type
    TC.checkSubExp [SubExp]
    vec
      Bool -> TypeM rep () -> TypeM rep ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    unless ([Type]
    vec_ts [Type] -> [Type] -> Bool
    forall a. Eq a => a -> a -> Bool
    == Lambda (Aliases rep) -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda (Aliases rep)
    lam) (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $
        ErrorCase rep -> TypeM rep ()
    forall rep a. ErrorCase rep -> TypeM rep a
    TC.bad (ErrorCase rep -> TypeM rep ())
    -> (Doc Any -> ErrorCase rep) -> Doc Any -> TypeM rep ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Text -> ErrorCase rep
    forall rep. Text -> ErrorCase rep
    TC.TypeError (Text -> ErrorCase rep)
    -> (Doc Any -> Text) -> Doc Any -> ErrorCase rep
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Doc Any -> Text
    forall a. Doc a -> Text
    docText (Doc Any -> TypeM rep ()) -> Doc Any -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $
          Doc Any
    "Return type"
            Doc Any -> Doc Any -> Doc Any
    forall a. Doc a -> Doc a -> Doc a
    </> Int -> Doc Any -> Doc Any
    forall ann. Int -> Doc ann -> Doc ann
    PP.indent Int
    2 ([Type] -> Doc Any
    forall ann. [Type] -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    pretty (Lambda (Aliases rep) -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda (Aliases rep)
    lam))
            Doc Any -> Doc Any -> Doc Any
    forall a. Doc a -> Doc a -> Doc a
    </> Doc Any
    "does not match type of seed vector"
            Doc Any -> Doc Any -> Doc Any
    forall a. Doc a -> Doc a -> Doc a
    </> Int -> Doc Any -> Doc Any
    forall ann. Int -> Doc ann -> Doc ann
    PP.indent Int
    2 ([Type] -> Doc Any
    forall ann. [Type] -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    pretty [Type]
    vec_ts)
    typeCheckSOAC (JVP [SubExp]
    args [SubExp]
    vec Lambda (Aliases rep)
    lam) = do
      [Arg]
    args' <- (SubExp -> TypeM rep Arg) -> [SubExp] -> TypeM rep [Arg]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM SubExp -> TypeM rep Arg
    forall rep. Checkable rep => SubExp -> TypeM rep Arg
    TC.checkArg [SubExp]
    args
      Lambda (Aliases rep) -> [Arg] -> TypeM rep ()
    forall rep.
    Checkable rep =>
    Lambda (Aliases rep) -> [Arg] -> TypeM rep ()
    TC.checkLambda Lambda (Aliases rep)
    lam ([Arg] -> TypeM rep ()) -> [Arg] -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $ (Arg -> Arg) -> [Arg] -> [Arg]
    forall a b. (a -> b) -> [a] -> [b]
    map Arg -> Arg
    TC.noArgAliases [Arg]
    args'
      [Type]
    vec_ts <- (SubExp -> TypeM rep Type) -> [SubExp] -> TypeM rep [Type]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM SubExp -> TypeM rep Type
    forall rep. Checkable rep => SubExp -> TypeM rep Type
    TC.checkSubExp [SubExp]
    vec
      Bool -> TypeM rep () -> TypeM rep ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    unless ([Type]
    vec_ts [Type] -> [Type] -> Bool
    forall a. Eq a => a -> a -> Bool
    == (Arg -> Type) -> [Arg] -> [Type]
    forall a b. (a -> b) -> [a] -> [b]
    map Arg -> Type
    TC.argType [Arg]
    args') (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $
        ErrorCase rep -> TypeM rep ()
    forall rep a. ErrorCase rep -> TypeM rep a
    TC.bad (ErrorCase rep -> TypeM rep ())
    -> (Doc Any -> ErrorCase rep) -> Doc Any -> TypeM rep ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Text -> ErrorCase rep
    forall rep. Text -> ErrorCase rep
    TC.TypeError (Text -> ErrorCase rep)
    -> (Doc Any -> Text) -> Doc Any -> ErrorCase rep
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Doc Any -> Text
    forall a. Doc a -> Text
    docText (Doc Any -> TypeM rep ()) -> Doc Any -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $
          Doc Any
    "Parameter type"
            Doc Any -> Doc Any -> Doc Any
    forall a. Doc a -> Doc a -> Doc a
    </> Int -> Doc Any -> Doc Any
    forall ann. Int -> Doc ann -> Doc ann
    PP.indent Int
    2 ([Type] -> Doc Any
    forall ann. [Type] -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    pretty ([Type] -> Doc Any) -> [Type] -> Doc Any
    forall a b. (a -> b) -> a -> b
    $ (Arg -> Type) -> [Arg] -> [Type]
    forall a b. (a -> b) -> [a] -> [b]
    map Arg -> Type
    TC.argType [Arg]
    args')
            Doc Any -> Doc Any -> Doc Any
    forall a. Doc a -> Doc a -> Doc a
    </> Doc Any
    "does not match type of seed vector"
            Doc Any -> Doc Any -> Doc Any
    forall a. Doc a -> Doc a -> Doc a
    </> Int -> Doc Any -> Doc Any
    forall ann. Int -> Doc ann -> Doc ann
    PP.indent Int
    2 ([Type] -> Doc Any
    forall ann. [Type] -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    pretty [Type]
    vec_ts)
    typeCheckSOAC (Stream SubExp
    size [VName]
    arrexps [SubExp]
    accexps Lambda (Aliases rep)
    lam) = do
      [Type] -> SubExp -> TypeM rep ()
    forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
    TC.require [PrimType -> Type
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    int64] SubExp
    size
      [Arg]
    accargs <- (SubExp -> TypeM rep Arg) -> [SubExp] -> TypeM rep [Arg]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM SubExp -> TypeM rep Arg
    forall rep. Checkable rep => SubExp -> TypeM rep Arg
    TC.checkArg [SubExp]
    accexps
      [Type]
    arrargs <- (VName -> TypeM rep Type) -> [VName] -> TypeM rep [Type]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM VName -> TypeM rep Type
    forall rep (m :: * -> *). HasScope rep m => VName -> m Type
    lookupType [VName]
    arrexps
      [Arg]
    _ <- SubExp -> [VName] -> TypeM rep [Arg]
    forall rep. Checkable rep => SubExp -> [VName] -> TypeM rep [Arg]
    TC.checkSOACArrayArgs SubExp
    size [VName]
    arrexps
      Param (LParamInfo rep)
    chunk <- case Lambda (Aliases rep) -> [LParam (Aliases rep)]
    forall {rep}. Lambda rep -> [Param (LParamInfo rep)]
    lambdaParams Lambda (Aliases rep)
    lam of
        LParam (Aliases rep)
    chunk : [LParam (Aliases rep)]
    _ -> Param (LParamInfo rep) -> TypeM rep (Param (LParamInfo rep))
    forall a. a -> TypeM rep a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure Param (LParamInfo rep)
    LParam (Aliases rep)
    chunk
        [] -> ErrorCase rep -> TypeM rep (Param (LParamInfo rep))
    forall rep a. ErrorCase rep -> TypeM rep a
    TC.bad (ErrorCase rep -> TypeM rep (Param (LParamInfo rep)))
    -> ErrorCase rep -> TypeM rep (Param (LParamInfo rep))
    forall a b. (a -> b) -> a -> b
    $ Text -> ErrorCase rep
    forall rep. Text -> ErrorCase rep
    TC.TypeError Text
    "Stream lambda without parameters."
      let asArg :: a -> (a, b)
    asArg a
    t = (a
    t, b
    forall a. Monoid a => a
    mempty)
          inttp :: TypeBase shape u
    inttp = PrimType -> TypeBase shape u
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    int64
          lamarrs' :: [Type]
    lamarrs' = (Type -> Type) -> [Type] -> [Type]
    forall a b. (a -> b) -> [a] -> [b]
    map (Type -> SubExp -> Type
    forall d u.
    ArrayShape (ShapeBase d) =>
    TypeBase (ShapeBase d) u -> d -> TypeBase (ShapeBase d) u
    `setOuterSize` VName -> SubExp
    Var (Param (LParamInfo rep) -> VName
    forall dec. Param dec -> VName
    paramName Param (LParamInfo rep)
    chunk)) [Type]
    arrargs
          acc_len :: Int
    acc_len = [SubExp] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [SubExp]
    accexps
          lamrtp :: [Type]
    lamrtp = Int -> [Type] -> [Type]
    forall a. Int -> [a] -> [a]
    take Int
    acc_len ([Type] -> [Type]) -> [Type] -> [Type]
    forall a b. (a -> b) -> a -> b
    $ Lambda (Aliases rep) -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda (Aliases rep)
    lam
      Bool -> TypeM rep () -> TypeM rep ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    unless ((Arg -> Type) -> [Arg] -> [Type]
    forall a b. (a -> b) -> [a] -> [b]
    map Arg -> Type
    TC.argType [Arg]
    accargs [Type] -> [Type] -> Bool
    forall a. Eq a => a -> a -> Bool
    == [Type]
    lamrtp) (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $
        ErrorCase rep -> TypeM rep ()
    forall rep a. ErrorCase rep -> TypeM rep a
    TC.bad (ErrorCase rep -> TypeM rep ())
    -> (Text -> ErrorCase rep) -> Text -> TypeM rep ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Text -> ErrorCase rep
    forall rep. Text -> ErrorCase rep
    TC.TypeError (Text -> TypeM rep ()) -> Text -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $
          Text
    "Stream with inconsistent accumulator type in lambda."
      -- just get the dflow of lambda on the fakearg, which does not alias
      -- arr, so we can later check that aliases of arr are not used inside lam.
      let fake_lamarrs' :: [Arg]
    fake_lamarrs' = (Type -> Arg) -> [Type] -> [Arg]
    forall a b. (a -> b) -> [a] -> [b]
    map Type -> Arg
    forall {b} {a}. Monoid b => a -> (a, b)
    asArg [Type]
    lamarrs'
      Lambda (Aliases rep) -> [Arg] -> TypeM rep ()
    forall rep.
    Checkable rep =>
    Lambda (Aliases rep) -> [Arg] -> TypeM rep ()
    TC.checkLambda Lambda (Aliases rep)
    lam ([Arg] -> TypeM rep ()) -> [Arg] -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $ Type -> Arg
    forall {b} {a}. Monoid b => a -> (a, b)
    asArg Type
    forall {shape} {u}. TypeBase shape u
    inttp Arg -> [Arg] -> [Arg]
    forall a. a -> [a] -> [a]
    : [Arg]
    accargs [Arg] -> [Arg] -> [Arg]
    forall a. [a] -> [a] -> [a]
    ++ [Arg]
    fake_lamarrs'
    typeCheckSOAC (Scatter SubExp
    w [VName]
    arrs ScatterSpec VName
    as Lambda (Aliases rep)
    lam) = do
      -- Requirements:
      --
      --   0. @lambdaReturnType@ of @lam@ must be a list
      --      [index types..., value types, ...].
      --
      --   1. The number of index types and value types must be equal to the number
      --      of return values from @lam@.
      --
      --   2. Each index type must have the type i64.
      --
      --   3. Each array in @as@ and the value types must have the same type
      --
      --   4. Each array in @as@ is consumed.  This is not really a check, but more
      --      of a requirement, so that e.g. the source is not hoisted out of a
      --      loop, which will mean it cannot be consumed.
      --
      --   5. Each of arrs must be an array matching a corresponding lambda
      --      parameters.
      --
      -- Code:

      -- First check the input size.
      [Type] -> SubExp -> TypeM rep ()
    forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
    TC.require [PrimType -> Type
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    int64] SubExp
    w

      -- 0.
      let ([Shape]
    as_ws, [Int]
    as_ns, [VName]
    _as_vs) = ScatterSpec VName -> ([Shape], [Int], [VName])
    forall a b c. [(a, b, c)] -> ([a], [b], [c])
    unzip3 ScatterSpec VName
    as
          indexes :: Int
    indexes = [Int] -> Int
    forall a. Num a => [a] -> a
    forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
    sum ([Int] -> Int) -> [Int] -> Int
    forall a b. (a -> b) -> a -> b
    $ (Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith Int -> Int -> Int
    forall a. Num a => a -> a -> a
    (*) [Int]
    as_ns ([Int] -> [Int]) -> [Int] -> [Int]
    forall a b. (a -> b) -> a -> b
    $ (Shape -> Int) -> [Shape] -> [Int]
    forall a b. (a -> b) -> [a] -> [b]
    map Shape -> Int
    forall a. ShapeBase a -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [Shape]
    as_ws
          rts :: [Type]
    rts = Lambda (Aliases rep) -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda (Aliases rep)
    lam
          rtsI :: [Type]
    rtsI = Int -> [Type] -> [Type]
    forall a. Int -> [a] -> [a]
    take Int
    indexes [Type]
    rts
          rtsV :: [Type]
    rtsV = Int -> [Type] -> [Type]
    forall a. Int -> [a] -> [a]
    drop Int
    indexes [Type]
    rts

      -- 1.
      Bool -> TypeM rep () -> TypeM rep ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    unless ([Type] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [Type]
    rts Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == [Int] -> Int
    forall a. Num a => [a] -> a
    forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
    sum [Int]
    as_ns Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + [Int] -> Int
    forall a. Num a => [a] -> a
    forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
    sum ((Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith Int -> Int -> Int
    forall a. Num a => a -> a -> a
    (*) [Int]
    as_ns ([Int] -> [Int]) -> [Int] -> [Int]
    forall a b. (a -> b) -> a -> b
    $ (Shape -> Int) -> [Shape] -> [Int]
    forall a b. (a -> b) -> [a] -> [b]
    map Shape -> Int
    forall a. ShapeBase a -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [Shape]
    as_ws)) (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $
        ErrorCase rep -> TypeM rep ()
    forall rep a. ErrorCase rep -> TypeM rep a
    TC.bad (ErrorCase rep -> TypeM rep ()) -> ErrorCase rep -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $
          Text -> ErrorCase rep
    forall rep. Text -> ErrorCase rep
    TC.TypeError Text
    "Scatter: number of index types, value types and array outputs do not match."

      -- 2.
      [Type] -> (Type -> TypeM rep ()) -> TypeM rep ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ [Type]
    rtsI ((Type -> TypeM rep ()) -> TypeM rep ())
    -> (Type -> TypeM rep ()) -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $ \Type
    rtI ->
        Bool -> TypeM rep () -> TypeM rep ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    unless (PrimType -> Type
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    int64 Type -> Type -> Bool
    forall a. Eq a => a -> a -> Bool
    == Type
    rtI) (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $
          ErrorCase rep -> TypeM rep ()
    forall rep a. ErrorCase rep -> TypeM rep a
    TC.bad (ErrorCase rep -> TypeM rep ()) -> ErrorCase rep -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $
            Text -> ErrorCase rep
    forall rep. Text -> ErrorCase rep
    TC.TypeError Text
    "Scatter: Index return type must be i64."

      [([Type], (Shape, Int, VName))]
    -> (([Type], (Shape, Int, VName)) -> TypeM rep ()) -> TypeM rep ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ ([[Type]] -> ScatterSpec VName -> [([Type], (Shape, Int, VName))]
    forall a b. [a] -> [b] -> [(a, b)]
    zip ([Int] -> [Type] -> [[Type]]
    forall a. [Int] -> [a] -> [[a]]
    chunks [Int]
    as_ns [Type]
    rtsV) ScatterSpec VName
    as) ((([Type], (Shape, Int, VName)) -> TypeM rep ()) -> TypeM rep ())
    -> (([Type], (Shape, Int, VName)) -> TypeM rep ()) -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $ \([Type]
    rtVs, (Shape
    aw, Int
    _, VName
    a)) -> do
        -- All lengths must have type i64.
        (SubExp -> TypeM rep ()) -> Shape -> TypeM rep ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ ([Type] -> SubExp -> TypeM rep ()
    forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
    TC.require [PrimType -> Type
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    int64]) Shape
    aw

        -- 3.
        [Type] -> (Type -> TypeM rep ()) -> TypeM rep ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ [Type]
    rtVs ((Type -> TypeM rep ()) -> TypeM rep ())
    -> (Type -> TypeM rep ()) -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $ \Type
    rtV -> [Type] -> VName -> TypeM rep ()
    forall rep. Checkable rep => [Type] -> VName -> TypeM rep ()
    TC.requireI [Type -> Shape -> Type
    arrayOfShape Type
    rtV Shape
    aw] VName
    a

        -- 4.
        Names -> TypeM rep ()
    forall rep. Checkable rep => Names -> TypeM rep ()
    TC.consume (Names -> TypeM rep ()) -> TypeM rep Names -> TypeM rep ()
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< VName -> TypeM rep Names
    forall rep. Checkable rep => VName -> TypeM rep Names
    TC.lookupAliases VName
    a

      -- 5.
      [Arg]
    arrargs <- SubExp -> [VName] -> TypeM rep [Arg]
    forall rep. Checkable rep => SubExp -> [VName] -> TypeM rep [Arg]
    TC.checkSOACArrayArgs SubExp
    w [VName]
    arrs
      Lambda (Aliases rep) -> [Arg] -> TypeM rep ()
    forall rep.
    Checkable rep =>
    Lambda (Aliases rep) -> [Arg] -> TypeM rep ()
    TC.checkLambda Lambda (Aliases rep)
    lam [Arg]
    arrargs
    typeCheckSOAC (Hist SubExp
    w [VName]
    arrs [HistOp (Aliases rep)]
    ops Lambda (Aliases rep)
    bucket_fun) = do
      [Type] -> SubExp -> TypeM rep ()
    forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
    TC.require [PrimType -> Type
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    int64] SubExp
    w

      -- Check the operators.
      [HistOp (Aliases rep)]
    -> (HistOp (Aliases rep) -> TypeM rep ()) -> TypeM rep ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ [HistOp (Aliases rep)]
    ops ((HistOp (Aliases rep) -> TypeM rep ()) -> TypeM rep ())
    -> (HistOp (Aliases rep) -> TypeM rep ()) -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $ \(HistOp Shape
    dest_shape SubExp
    rf [VName]
    dests [SubExp]
    nes Lambda (Aliases rep)
    op) -> do
        [Arg]
    nes' <- (SubExp -> TypeM rep Arg) -> [SubExp] -> TypeM rep [Arg]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM SubExp -> TypeM rep Arg
    forall rep. Checkable rep => SubExp -> TypeM rep Arg
    TC.checkArg [SubExp]
    nes
        (SubExp -> TypeM rep ()) -> Shape -> TypeM rep ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ ([Type] -> SubExp -> TypeM rep ()
    forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
    TC.require [PrimType -> Type
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    int64]) Shape
    dest_shape
        [Type] -> SubExp -> TypeM rep ()
    forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
    TC.require [PrimType -> Type
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    int64] SubExp
    rf

        -- Operator type must match the type of neutral elements.
        Lambda (Aliases rep) -> [Arg] -> TypeM rep ()
    forall rep.
    Checkable rep =>
    Lambda (Aliases rep) -> [Arg] -> TypeM rep ()
    TC.checkLambda Lambda (Aliases rep)
    op ([Arg] -> TypeM rep ()) -> [Arg] -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $ (Arg -> Arg) -> [Arg] -> [Arg]
    forall a b. (a -> b) -> [a] -> [b]
    map Arg -> Arg
    TC.noArgAliases ([Arg] -> [Arg]) -> [Arg] -> [Arg]
    forall a b. (a -> b) -> a -> b
    $ [Arg]
    nes' [Arg] -> [Arg] -> [Arg]
    forall a. [a] -> [a] -> [a]
    ++ [Arg]
    nes'
        let nes_t :: [Type]
    nes_t = (Arg -> Type) -> [Arg] -> [Type]
    forall a b. (a -> b) -> [a] -> [b]
    map Arg -> Type
    TC.argType [Arg]
    nes'
        Bool -> TypeM rep () -> TypeM rep ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    unless ([Type]
    nes_t [Type] -> [Type] -> Bool
    forall a. Eq a => a -> a -> Bool
    == Lambda (Aliases rep) -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda (Aliases rep)
    op) (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $
          ErrorCase rep -> TypeM rep ()
    forall rep a. ErrorCase rep -> TypeM rep a
    TC.bad (ErrorCase rep -> TypeM rep ())
    -> (Text -> ErrorCase rep) -> Text -> TypeM rep ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Text -> ErrorCase rep
    forall rep. Text -> ErrorCase rep
    TC.TypeError (Text -> TypeM rep ()) -> Text -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $
            Text
    "Operator has return type "
              Text -> Text -> Text
    forall a. Semigroup a => a -> a -> a
    <> [Type] -> Text
    forall a. Pretty a => [a] -> Text
    prettyTuple (Lambda (Aliases rep) -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda (Aliases rep)
    op)
              Text -> Text -> Text
    forall a. Semigroup a => a -> a -> a
    <> Text
    " but neutral element has type "
              Text -> Text -> Text
    forall a. Semigroup a => a -> a -> a
    <> [Type] -> Text
    forall a. Pretty a => [a] -> Text
    prettyTuple [Type]
    nes_t

        -- Arrays must have proper type.
        [(Type, VName)] -> ((Type, VName) -> TypeM rep ()) -> TypeM rep ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    t a -> (a -> m b) -> m ()
    forM_ ([Type] -> [VName] -> [(Type, VName)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [Type]
    nes_t [VName]
    dests) (((Type, VName) -> TypeM rep ()) -> TypeM rep ())
    -> ((Type, VName) -> TypeM rep ()) -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $ \(Type
    t, VName
    dest) -> do
          [Type] -> VName -> TypeM rep ()
    forall rep. Checkable rep => [Type] -> VName -> TypeM rep ()
    TC.requireI [Type
    t Type -> Shape -> Type
    `arrayOfShape` Shape
    dest_shape] VName
    dest
          Names -> TypeM rep ()
    forall rep. Checkable rep => Names -> TypeM rep ()
    TC.consume (Names -> TypeM rep ()) -> TypeM rep Names -> TypeM rep ()
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
    =<< VName -> TypeM rep Names
    forall rep. Checkable rep => VName -> TypeM rep Names
    TC.lookupAliases VName
    dest

      -- Types of input arrays must equal parameter types for bucket function.
      [Arg]
    img' <- SubExp -> [VName] -> TypeM rep [Arg]
    forall rep. Checkable rep => SubExp -> [VName] -> TypeM rep [Arg]
    TC.checkSOACArrayArgs SubExp
    w [VName]
    arrs
      Lambda (Aliases rep) -> [Arg] -> TypeM rep ()
    forall rep.
    Checkable rep =>
    Lambda (Aliases rep) -> [Arg] -> TypeM rep ()
    TC.checkLambda Lambda (Aliases rep)
    bucket_fun [Arg]
    img'

      -- Return type of bucket function must be an index for each
      -- operation followed by the values to write.
      [Type]
    nes_ts <- [[Type]] -> [Type]
    forall (t :: * -> *) a. Foldable t => t [a] -> [a]
    concat ([[Type]] -> [Type]) -> TypeM rep [[Type]] -> TypeM rep [Type]
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (HistOp (Aliases rep) -> TypeM rep [Type])
    -> [HistOp (Aliases rep)] -> TypeM rep [[Type]]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM ((SubExp -> TypeM rep Type) -> [SubExp] -> TypeM rep [Type]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM SubExp -> TypeM rep Type
    forall t (m :: * -> *). HasScope t m => SubExp -> m Type
    subExpType ([SubExp] -> TypeM rep [Type])
    -> (HistOp (Aliases rep) -> [SubExp])
    -> HistOp (Aliases rep)
    -> TypeM rep [Type]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . HistOp (Aliases rep) -> [SubExp]
    forall rep. HistOp rep -> [SubExp]
    histNeutral) [HistOp (Aliases rep)]
    ops
      let bucket_ret_t :: [Type]
    bucket_ret_t =
            (HistOp (Aliases rep) -> [Type])
    -> [HistOp (Aliases rep)] -> [Type]
    forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
    concatMap ((Int -> Type -> [Type]
    forall a. Int -> a -> [a]
    `replicate` PrimType -> Type
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    int64) (Int -> [Type])
    -> (HistOp (Aliases rep) -> Int) -> HistOp (Aliases rep) -> [Type]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Shape -> Int
    forall a. ArrayShape a => a -> Int
    shapeRank (Shape -> Int)
    -> (HistOp (Aliases rep) -> Shape) -> HistOp (Aliases rep) -> Int
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . HistOp (Aliases rep) -> Shape
    forall rep. HistOp rep -> Shape
    histShape) [HistOp (Aliases rep)]
    ops
              [Type] -> [Type] -> [Type]
    forall a. [a] -> [a] -> [a]
    ++ [Type]
    nes_ts
      Bool -> TypeM rep () -> TypeM rep ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    unless ([Type]
    bucket_ret_t [Type] -> [Type] -> Bool
    forall a. Eq a => a -> a -> Bool
    == Lambda (Aliases rep) -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda (Aliases rep)
    bucket_fun) (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $
        ErrorCase rep -> TypeM rep ()
    forall rep a. ErrorCase rep -> TypeM rep a
    TC.bad (ErrorCase rep -> TypeM rep ())
    -> (Text -> ErrorCase rep) -> Text -> TypeM rep ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Text -> ErrorCase rep
    forall rep. Text -> ErrorCase rep
    TC.TypeError (Text -> TypeM rep ()) -> Text -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $
          Text
    "Bucket function has return type "
            Text -> Text -> Text
    forall a. Semigroup a => a -> a -> a
    <> [Type] -> Text
    forall a. Pretty a => [a] -> Text
    prettyTuple (Lambda (Aliases rep) -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda (Aliases rep)
    bucket_fun)
            Text -> Text -> Text
    forall a. Semigroup a => a -> a -> a
    <> Text
    " but should have type "
            Text -> Text -> Text
    forall a. Semigroup a => a -> a -> a
    <> [Type] -> Text
    forall a. Pretty a => [a] -> Text
    prettyTuple [Type]
    bucket_ret_t
    typeCheckSOAC (Screma SubExp
    w [VName]
    arrs (ScremaForm Lambda (Aliases rep)
    map_lam [Scan (Aliases rep)]
    scans [Reduce (Aliases rep)]
    reds)) = do
      [Type] -> SubExp -> TypeM rep ()
    forall rep. Checkable rep => [Type] -> SubExp -> TypeM rep ()
    TC.require [PrimType -> Type
    forall shape u. PrimType -> TypeBase shape u
    Prim PrimType
    int64] SubExp
    w
      [Arg]
    arrs' <- SubExp -> [VName] -> TypeM rep [Arg]
    forall rep. Checkable rep => SubExp -> [VName] -> TypeM rep [Arg]
    TC.checkSOACArrayArgs SubExp
    w [VName]
    arrs
      Lambda (Aliases rep) -> [Arg] -> TypeM rep ()
    forall rep.
    Checkable rep =>
    Lambda (Aliases rep) -> [Arg] -> TypeM rep ()
    TC.checkLambda Lambda (Aliases rep)
    map_lam [Arg]
    arrs'

      [Arg]
    scan_nes' <- ([[Arg]] -> [Arg]) -> TypeM rep [[Arg]] -> TypeM rep [Arg]
    forall a b. (a -> b) -> TypeM rep a -> TypeM rep b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap [[Arg]] -> [Arg]
    forall (t :: * -> *) a. Foldable t => t [a] -> [a]
    concat (TypeM rep [[Arg]] -> TypeM rep [Arg])
    -> TypeM rep [[Arg]] -> TypeM rep [Arg]
    forall a b. (a -> b) -> a -> b
    $
        [Scan (Aliases rep)]
    -> (Scan (Aliases rep) -> TypeM rep [Arg]) -> TypeM rep [[Arg]]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM [Scan (Aliases rep)]
    scans ((Scan (Aliases rep) -> TypeM rep [Arg]) -> TypeM rep [[Arg]])
    -> (Scan (Aliases rep) -> TypeM rep [Arg]) -> TypeM rep [[Arg]]
    forall a b. (a -> b) -> a -> b
    $ \(Scan Lambda (Aliases rep)
    scan_lam [SubExp]
    scan_nes) -> do
          [Arg]
    scan_nes' <- (SubExp -> TypeM rep Arg) -> [SubExp] -> TypeM rep [Arg]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM SubExp -> TypeM rep Arg
    forall rep. Checkable rep => SubExp -> TypeM rep Arg
    TC.checkArg [SubExp]
    scan_nes
          let scan_t :: [Type]
    scan_t = (Arg -> Type) -> [Arg] -> [Type]
    forall a b. (a -> b) -> [a] -> [b]
    map Arg -> Type
    TC.argType [Arg]
    scan_nes'
          Lambda (Aliases rep) -> [Arg] -> TypeM rep ()
    forall rep.
    Checkable rep =>
    Lambda (Aliases rep) -> [Arg] -> TypeM rep ()
    TC.checkLambda Lambda (Aliases rep)
    scan_lam ([Arg] -> TypeM rep ()) -> [Arg] -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $ (Arg -> Arg) -> [Arg] -> [Arg]
    forall a b. (a -> b) -> [a] -> [b]
    map Arg -> Arg
    TC.noArgAliases ([Arg] -> [Arg]) -> [Arg] -> [Arg]
    forall a b. (a -> b) -> a -> b
    $ [Arg]
    scan_nes' [Arg] -> [Arg] -> [Arg]
    forall a. [a] -> [a] -> [a]
    ++ [Arg]
    scan_nes'
          Bool -> TypeM rep () -> TypeM rep ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    unless ([Type]
    scan_t [Type] -> [Type] -> Bool
    forall a. Eq a => a -> a -> Bool
    == Lambda (Aliases rep) -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda (Aliases rep)
    scan_lam) (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $
            ErrorCase rep -> TypeM rep ()
    forall rep a. ErrorCase rep -> TypeM rep a
    TC.bad (ErrorCase rep -> TypeM rep ())
    -> (Text -> ErrorCase rep) -> Text -> TypeM rep ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Text -> ErrorCase rep
    forall rep. Text -> ErrorCase rep
    TC.TypeError (Text -> TypeM rep ()) -> Text -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $
              Text
    "Scan function returns type "
                Text -> Text -> Text
    forall a. Semigroup a => a -> a -> a
    <> [Type] -> Text
    forall a. Pretty a => [a] -> Text
    prettyTuple (Lambda (Aliases rep) -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda (Aliases rep)
    scan_lam)
                Text -> Text -> Text
    forall a. Semigroup a => a -> a -> a
    <> Text
    " but neutral element has type "
                Text -> Text -> Text
    forall a. Semigroup a => a -> a -> a
    <> [Type] -> Text
    forall a. Pretty a => [a] -> Text
    prettyTuple [Type]
    scan_t
          [Arg] -> TypeM rep [Arg]
    forall a. a -> TypeM rep a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [Arg]
    scan_nes'

      [Arg]
    red_nes' <- ([[Arg]] -> [Arg]) -> TypeM rep [[Arg]] -> TypeM rep [Arg]
    forall a b. (a -> b) -> TypeM rep a -> TypeM rep b
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    fmap [[Arg]] -> [Arg]
    forall (t :: * -> *) a. Foldable t => t [a] -> [a]
    concat (TypeM rep [[Arg]] -> TypeM rep [Arg])
    -> TypeM rep [[Arg]] -> TypeM rep [Arg]
    forall a b. (a -> b) -> a -> b
    $
        [Reduce (Aliases rep)]
    -> (Reduce (Aliases rep) -> TypeM rep [Arg]) -> TypeM rep [[Arg]]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    t a -> (a -> m b) -> m (t b)
    forM [Reduce (Aliases rep)]
    reds ((Reduce (Aliases rep) -> TypeM rep [Arg]) -> TypeM rep [[Arg]])
    -> (Reduce (Aliases rep) -> TypeM rep [Arg]) -> TypeM rep [[Arg]]
    forall a b. (a -> b) -> a -> b
    $ \(Reduce Commutativity
    _ Lambda (Aliases rep)
    red_lam [SubExp]
    red_nes) -> do
          [Arg]
    red_nes' <- (SubExp -> TypeM rep Arg) -> [SubExp] -> TypeM rep [Arg]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM SubExp -> TypeM rep Arg
    forall rep. Checkable rep => SubExp -> TypeM rep Arg
    TC.checkArg [SubExp]
    red_nes
          let red_t :: [Type]
    red_t = (Arg -> Type) -> [Arg] -> [Type]
    forall a b. (a -> b) -> [a] -> [b]
    map Arg -> Type
    TC.argType [Arg]
    red_nes'
          Lambda (Aliases rep) -> [Arg] -> TypeM rep ()
    forall rep.
    Checkable rep =>
    Lambda (Aliases rep) -> [Arg] -> TypeM rep ()
    TC.checkLambda Lambda (Aliases rep)
    red_lam ([Arg] -> TypeM rep ()) -> [Arg] -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $ (Arg -> Arg) -> [Arg] -> [Arg]
    forall a b. (a -> b) -> [a] -> [b]
    map Arg -> Arg
    TC.noArgAliases ([Arg] -> [Arg]) -> [Arg] -> [Arg]
    forall a b. (a -> b) -> a -> b
    $ [Arg]
    red_nes' [Arg] -> [Arg] -> [Arg]
    forall a. [a] -> [a] -> [a]
    ++ [Arg]
    red_nes'
          Bool -> TypeM rep () -> TypeM rep ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    unless ([Type]
    red_t [Type] -> [Type] -> Bool
    forall a. Eq a => a -> a -> Bool
    == Lambda (Aliases rep) -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda (Aliases rep)
    red_lam) (TypeM rep () -> TypeM rep ()) -> TypeM rep () -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $
            ErrorCase rep -> TypeM rep ()
    forall rep a. ErrorCase rep -> TypeM rep a
    TC.bad (ErrorCase rep -> TypeM rep ())
    -> (Text -> ErrorCase rep) -> Text -> TypeM rep ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Text -> ErrorCase rep
    forall rep. Text -> ErrorCase rep
    TC.TypeError (Text -> TypeM rep ()) -> Text -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $
              Text
    "Reduce function returns type "
                Text -> Text -> Text
    forall a. Semigroup a => a -> a -> a
    <> [Type] -> Text
    forall a. Pretty a => [a] -> Text
    prettyTuple (Lambda (Aliases rep) -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda (Aliases rep)
    red_lam)
                Text -> Text -> Text
    forall a. Semigroup a => a -> a -> a
    <> Text
    " but neutral element has type "
                Text -> Text -> Text
    forall a. Semigroup a => a -> a -> a
    <> [Type] -> Text
    forall a. Pretty a => [a] -> Text
    prettyTuple [Type]
    red_t
          [Arg] -> TypeM rep [Arg]
    forall a. a -> TypeM rep a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [Arg]
    red_nes'

      let map_lam_ts :: [Type]
    map_lam_ts = Lambda (Aliases rep) -> [Type]
    forall rep. Lambda rep -> [Type]
    lambdaReturnType Lambda (Aliases rep)
    map_lam

      Bool -> TypeM rep () -> TypeM rep ()
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
    unless
        ( Int -> [Type] -> [Type]
    forall a. Int -> [a] -> [a]
    take ([Arg] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [Arg]
    scan_nes' Int -> Int -> Int
    forall a. Num a => a -> a -> a
    + [Arg] -> Int
    forall a. [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [Arg]
    red_nes') [Type]
    map_lam_ts
            [Type] -> [Type] -> Bool
    forall a. Eq a => a -> a -> Bool
    == (Arg -> Type) -> [Arg] -> [Type]
    forall a b. (a -> b) -> [a] -> [b]
    map Arg -> Type
    TC.argType ([Arg]
    scan_nes' [Arg] -> [Arg] -> [Arg]
    forall a. [a] -> [a] -> [a]
    ++ [Arg]
    red_nes')
        )
        (TypeM rep () -> TypeM rep ())
    -> (Text -> TypeM rep ()) -> Text -> TypeM rep ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . ErrorCase rep -> TypeM rep ()
    forall rep a. ErrorCase rep -> TypeM rep a
    TC.bad
        (ErrorCase rep -> TypeM rep ())
    -> (Text -> ErrorCase rep) -> Text -> TypeM rep ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Text -> ErrorCase rep
    forall rep. Text -> ErrorCase rep
    TC.TypeError
        (Text -> TypeM rep ()) -> Text -> TypeM rep ()
    forall a b. (a -> b) -> a -> b
    $ Text
    "Map function return type "
          Text -> Text -> Text
    forall a. Semigroup a => a -> a -> a
    <> [Type] -> Text
    forall a. Pretty a => [a] -> Text
    prettyTuple [Type]
    map_lam_ts
          Text -> Text -> Text
    forall a. Semigroup a => a -> a -> a
    <> Text
    " wrong for given scan and reduction functions."

    instance RephraseOp SOAC where
      rephraseInOp :: forall (m :: * -> *) from to.
    Monad m =>
    Rephraser m from to -> SOAC from -> m (SOAC to)
    rephraseInOp Rephraser m from to
    r (VJP [SubExp]
    args [SubExp]
    vec Lambda from
    lam) =
        [SubExp] -> [SubExp] -> Lambda to -> SOAC to
    forall rep. [SubExp] -> [SubExp] -> Lambda rep -> SOAC rep
    VJP [SubExp]
    args [SubExp]
    vec (Lambda to -> SOAC to) -> m (Lambda to) -> m (SOAC to)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Rephraser m from to -> Lambda from -> m (Lambda to)
    forall (m :: * -> *) from to.
    Monad m =>
    Rephraser m from to -> Lambda from -> m (Lambda to)
    rephraseLambda Rephraser m from to
    r Lambda from
    lam
      rephraseInOp Rephraser m from to
    r (JVP [SubExp]
    args [SubExp]
    vec Lambda from
    lam) =
        [SubExp] -> [SubExp] -> Lambda to -> SOAC to
    forall rep. [SubExp] -> [SubExp] -> Lambda rep -> SOAC rep
    JVP [SubExp]
    args [SubExp]
    vec (Lambda to -> SOAC to) -> m (Lambda to) -> m (SOAC to)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Rephraser m from to -> Lambda from -> m (Lambda to)
    forall (m :: * -> *) from to.
    Monad m =>
    Rephraser m from to -> Lambda from -> m (Lambda to)
    rephraseLambda Rephraser m from to
    r Lambda from
    lam
      rephraseInOp Rephraser m from to
    r (Stream SubExp
    w [VName]
    arrs [SubExp]
    acc Lambda from
    lam) =
        SubExp -> [VName] -> [SubExp] -> Lambda to -> SOAC to
    forall rep. SubExp -> [VName] -> [SubExp] -> Lambda rep -> SOAC rep
    Stream SubExp
    w [VName]
    arrs [SubExp]
    acc (Lambda to -> SOAC to) -> m (Lambda to) -> m (SOAC to)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Rephraser m from to -> Lambda from -> m (Lambda to)
    forall (m :: * -> *) from to.
    Monad m =>
    Rephraser m from to -> Lambda from -> m (Lambda to)
    rephraseLambda Rephraser m from to
    r Lambda from
    lam
      rephraseInOp Rephraser m from to
    r (Scatter SubExp
    w [VName]
    arrs ScatterSpec VName
    dests Lambda from
    lam) =
        SubExp -> [VName] -> ScatterSpec VName -> Lambda to -> SOAC to
    forall rep.
    SubExp -> [VName] -> ScatterSpec VName -> Lambda rep -> SOAC rep
    Scatter SubExp
    w [VName]
    arrs ScatterSpec VName
    dests (Lambda to -> SOAC to) -> m (Lambda to) -> m (SOAC to)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Rephraser m from to -> Lambda from -> m (Lambda to)
    forall (m :: * -> *) from to.
    Monad m =>
    Rephraser m from to -> Lambda from -> m (Lambda to)
    rephraseLambda Rephraser m from to
    r Lambda from
    lam
      rephraseInOp Rephraser m from to
    r (Hist SubExp
    w [VName]
    arrs [HistOp from]
    ops Lambda from
    lam) =
        SubExp -> [VName] -> [HistOp to] -> Lambda to -> SOAC to
    forall rep.
    SubExp -> [VName] -> [HistOp rep] -> Lambda rep -> SOAC rep
    Hist SubExp
    w [VName]
    arrs ([HistOp to] -> Lambda to -> SOAC to)
    -> m [HistOp to] -> m (Lambda to -> SOAC to)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> (HistOp from -> m (HistOp to)) -> [HistOp from] -> m [HistOp to]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM HistOp from -> m (HistOp to)
    onOp [HistOp from]
    ops m (Lambda to -> SOAC to) -> m (Lambda to) -> m (SOAC to)
    forall a b. m (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> Rephraser m from to -> Lambda from -> m (Lambda to)
    forall (m :: * -> *) from to.
    Monad m =>
    Rephraser m from to -> Lambda from -> m (Lambda to)
    rephraseLambda Rephraser m from to
    r Lambda from
    lam
        where
          onOp :: HistOp from -> m (HistOp to)
    onOp (HistOp Shape
    dest_shape SubExp
    rf [VName]
    dests [SubExp]
    nes Lambda from
    op) =
            Shape -> SubExp -> [VName] -> [SubExp] -> Lambda to -> HistOp to
    forall rep.
    Shape -> SubExp -> [VName] -> [SubExp] -> Lambda rep -> HistOp rep
    HistOp Shape
    dest_shape SubExp
    rf [VName]
    dests [SubExp]
    nes (Lambda to -> HistOp to) -> m (Lambda to) -> m (HistOp to)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Rephraser m from to -> Lambda from -> m (Lambda to)
    forall (m :: * -> *) from to.
    Monad m =>
    Rephraser m from to -> Lambda from -> m (Lambda to)
    rephraseLambda Rephraser m from to
    r Lambda from
    op
      rephraseInOp Rephraser m from to
    r (Screma SubExp
    w [VName]
    arrs (ScremaForm Lambda from
    lam [Scan from]
    scans [Reduce from]
    red)) =
        SubExp -> [VName] -> ScremaForm to -> SOAC to
    forall rep. SubExp -> [VName] -> ScremaForm rep -> SOAC rep
    Screma SubExp
    w [VName]
    arrs
          (ScremaForm to -> SOAC to) -> m (ScremaForm to) -> m (SOAC to)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> ( Lambda to -> [Scan to] -> [Reduce to] -> ScremaForm to
    forall rep.
    Lambda rep -> [Scan rep] -> [Reduce rep] -> ScremaForm rep
    ScremaForm
                  (Lambda to -> [Scan to] -> [Reduce to] -> ScremaForm to)
    -> m (Lambda to) -> m ([Scan to] -> [Reduce to] -> ScremaForm to)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Rephraser m from to -> Lambda from -> m (Lambda to)
    forall (m :: * -> *) from to.
    Monad m =>
    Rephraser m from to -> Lambda from -> m (Lambda to)
    rephraseLambda Rephraser m from to
    r Lambda from
    lam
                  m ([Scan to] -> [Reduce to] -> ScremaForm to)
    -> m [Scan to] -> m ([Reduce to] -> ScremaForm to)
    forall a b. m (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> (Scan from -> m (Scan to)) -> [Scan from] -> m [Scan to]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM Scan from -> m (Scan to)
    onScan [Scan from]
    scans
                  m ([Reduce to] -> ScremaForm to)
    -> m [Reduce to] -> m (ScremaForm to)
    forall a b. m (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> (Reduce from -> m (Reduce to)) -> [Reduce from] -> m [Reduce to]
    forall (t :: * -> *) (m :: * -> *) a b.
    (Traversable t, Monad m) =>
    (a -> m b) -> t a -> m (t b)
    forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
    mapM Reduce from -> m (Reduce to)
    onRed [Reduce from]
    red
              )
        where
          onScan :: Scan from -> m (Scan to)
    onScan (Scan Lambda from
    op [SubExp]
    nes) = Lambda to -> [SubExp] -> Scan to
    forall rep. Lambda rep -> [SubExp] -> Scan rep
    Scan (Lambda to -> [SubExp] -> Scan to)
    -> m (Lambda to) -> m ([SubExp] -> Scan to)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Rephraser m from to -> Lambda from -> m (Lambda to)
    forall (m :: * -> *) from to.
    Monad m =>
    Rephraser m from to -> Lambda from -> m (Lambda to)
    rephraseLambda Rephraser m from to
    r Lambda from
    op m ([SubExp] -> Scan to) -> m [SubExp] -> m (Scan to)
    forall a b. m (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> [SubExp] -> m [SubExp]
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [SubExp]
    nes
          onRed :: Reduce from -> m (Reduce to)
    onRed (Reduce Commutativity
    comm Lambda from
    op [SubExp]
    nes) = Commutativity -> Lambda to -> [SubExp] -> Reduce to
    forall rep. Commutativity -> Lambda rep -> [SubExp] -> Reduce rep
    Reduce Commutativity
    comm (Lambda to -> [SubExp] -> Reduce to)
    -> m (Lambda to) -> m ([SubExp] -> Reduce to)
    forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
    <$> Rephraser m from to -> Lambda from -> m (Lambda to)
    forall (m :: * -> *) from to.
    Monad m =>
    Rephraser m from to -> Lambda from -> m (Lambda to)
    rephraseLambda Rephraser m from to
    r Lambda from
    op m ([SubExp] -> Reduce to) -> m [SubExp] -> m (Reduce to)
    forall a b. m (a -> b) -> m a -> m b
    forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
    <*> [SubExp] -> m [SubExp]
    forall a. a -> m a
    forall (f :: * -> *) a. Applicative f => a -> f a
    pure [SubExp]
    nes

    instance (OpMetrics (Op rep)) => OpMetrics (SOAC rep) where
      opMetrics :: SOAC rep -> MetricsM ()
    opMetrics (VJP [SubExp]
    _ [SubExp]
    _ Lambda rep
    lam) =
        Text -> MetricsM () -> MetricsM ()
    inside Text
    "VJP" (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
    forall a b. (a -> b) -> a -> b
    $ Lambda rep -> MetricsM ()
    forall rep. OpMetrics (Op rep) => Lambda rep -> MetricsM ()
    lambdaMetrics Lambda rep
    lam
      opMetrics (JVP [SubExp]
    _ [SubExp]
    _ Lambda rep
    lam) =
        Text -> MetricsM () -> MetricsM ()
    inside Text
    "JVP" (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
    forall a b. (a -> b) -> a -> b
    $ Lambda rep -> MetricsM ()
    forall rep. OpMetrics (Op rep) => Lambda rep -> MetricsM ()
    lambdaMetrics Lambda rep
    lam
      opMetrics (Stream SubExp
    _ [VName]
    _ [SubExp]
    _ Lambda rep
    lam) =
        Text -> MetricsM () -> MetricsM ()
    inside Text
    "Stream" (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
    forall a b. (a -> b) -> a -> b
    $ Lambda rep -> MetricsM ()
    forall rep. OpMetrics (Op rep) => Lambda rep -> MetricsM ()
    lambdaMetrics Lambda rep
    lam
      opMetrics (Scatter SubExp
    _len [VName]
    _ ScatterSpec VName
    _ Lambda rep
    lam) =
        Text -> MetricsM () -> MetricsM ()
    inside Text
    "Scatter" (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
    forall a b. (a -> b) -> a -> b
    $ Lambda rep -> MetricsM ()
    forall rep. OpMetrics (Op rep) => Lambda rep -> MetricsM ()
    lambdaMetrics Lambda rep
    lam
      opMetrics (Hist SubExp
    _ [VName]
    _ [HistOp rep]
    ops Lambda rep
    bucket_fun) =
        Text -> MetricsM () -> MetricsM ()
    inside Text
    "Hist" (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
    forall a b. (a -> b) -> a -> b
    $ (HistOp rep -> MetricsM ()) -> [HistOp rep] -> MetricsM ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ (Lambda rep -> MetricsM ()
    forall rep. OpMetrics (Op rep) => Lambda rep -> MetricsM ()
    lambdaMetrics (Lambda rep -> MetricsM ())
    -> (HistOp rep -> Lambda rep) -> HistOp rep -> MetricsM ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . HistOp rep -> Lambda rep
    forall rep. HistOp rep -> Lambda rep
    histOp) [HistOp rep]
    ops MetricsM () -> MetricsM () -> MetricsM ()
    forall a b. MetricsM a -> MetricsM b -> MetricsM b
    forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
    >> Lambda rep -> MetricsM ()
    forall rep. OpMetrics (Op rep) => Lambda rep -> MetricsM ()
    lambdaMetrics Lambda rep
    bucket_fun
      opMetrics (Screma SubExp
    _ [VName]
    _ (ScremaForm Lambda rep
    map_lam [Scan rep]
    scans [Reduce rep]
    reds)) =
        Text -> MetricsM () -> MetricsM ()
    inside Text
    "Screma" (MetricsM () -> MetricsM ()) -> MetricsM () -> MetricsM ()
    forall a b. (a -> b) -> a -> b
    $ do
          Lambda rep -> MetricsM ()
    forall rep. OpMetrics (Op rep) => Lambda rep -> MetricsM ()
    lambdaMetrics Lambda rep
    map_lam
          (Scan rep -> MetricsM ()) -> [Scan rep] -> MetricsM ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ (Lambda rep -> MetricsM ()
    forall rep. OpMetrics (Op rep) => Lambda rep -> MetricsM ()
    lambdaMetrics (Lambda rep -> MetricsM ())
    -> (Scan rep -> Lambda rep) -> Scan rep -> MetricsM ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Scan rep -> Lambda rep
    forall rep. Scan rep -> Lambda rep
    scanLambda) [Scan rep]
    scans
          (Reduce rep -> MetricsM ()) -> [Reduce rep] -> MetricsM ()
    forall (t :: * -> *) (m :: * -> *) a b.
    (Foldable t, Monad m) =>
    (a -> m b) -> t a -> m ()
    mapM_ (Lambda rep -> MetricsM ()
    forall rep. OpMetrics (Op rep) => Lambda rep -> MetricsM ()
    lambdaMetrics (Lambda rep -> MetricsM ())
    -> (Reduce rep -> Lambda rep) -> Reduce rep -> MetricsM ()
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Reduce rep -> Lambda rep
    forall rep. Reduce rep -> Lambda rep
    redLambda) [Reduce rep]
    reds

    instance (PrettyRep rep) => PP.Pretty (SOAC rep) where
      pretty :: forall ann. SOAC rep -> Doc ann
    pretty (VJP [SubExp]
    args [SubExp]
    vec Lambda rep
    lam) =
        Doc ann
    "vjp"
          Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    parens
            ( Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    PP.align (Doc ann -> Doc ann) -> Doc ann -> Doc ann
    forall a b. (a -> b) -> a -> b
    $
                Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    PP.braces ([Doc ann] -> Doc ann
    forall a. [Doc a] -> Doc a
    commasep ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
    forall a b. (a -> b) -> a -> b
    $ (SubExp -> Doc ann) -> [SubExp] -> [Doc ann]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExp -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. SubExp -> Doc ann
    pretty [SubExp]
    args)
                  Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    comma Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a
    </> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    PP.braces ([Doc ann] -> Doc ann
    forall a. [Doc a] -> Doc a
    commasep ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
    forall a b. (a -> b) -> a -> b
    $ (SubExp -> Doc ann) -> [SubExp] -> [Doc ann]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExp -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. SubExp -> Doc ann
    pretty [SubExp]
    vec)
                  Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    comma Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a
    </> Lambda rep -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. Lambda rep -> Doc ann
    pretty Lambda rep
    lam
            )
      pretty (JVP [SubExp]
    args [SubExp]
    vec Lambda rep
    lam) =
        Doc ann
    "jvp"
          Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    parens
            ( Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    PP.align (Doc ann -> Doc ann) -> Doc ann -> Doc ann
    forall a b. (a -> b) -> a -> b
    $
                Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    PP.braces ([Doc ann] -> Doc ann
    forall a. [Doc a] -> Doc a
    commasep ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
    forall a b. (a -> b) -> a -> b
    $ (SubExp -> Doc ann) -> [SubExp] -> [Doc ann]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExp -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. SubExp -> Doc ann
    pretty [SubExp]
    args)
                  Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    comma Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a
    </> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    PP.braces ([Doc ann] -> Doc ann
    forall a. [Doc a] -> Doc a
    commasep ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
    forall a b. (a -> b) -> a -> b
    $ (SubExp -> Doc ann) -> [SubExp] -> [Doc ann]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExp -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. SubExp -> Doc ann
    pretty [SubExp]
    vec)
                  Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    comma Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a
    </> Lambda rep -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. Lambda rep -> Doc ann
    pretty Lambda rep
    lam
            )
      pretty (Stream SubExp
    size [VName]
    arrs [SubExp]
    acc Lambda rep
    lam) =
        SubExp -> [VName] -> [SubExp] -> Lambda rep -> Doc ann
    forall rep inp ann.
    (PrettyRep rep, Pretty inp) =>
    SubExp -> [inp] -> [SubExp] -> Lambda rep -> Doc ann
    ppStream SubExp
    size [VName]
    arrs [SubExp]
    acc Lambda rep
    lam
      pretty (Scatter SubExp
    w [VName]
    arrs ScatterSpec VName
    dests Lambda rep
    lam) =
        SubExp -> [VName] -> ScatterSpec VName -> Lambda rep -> Doc ann
    forall rep inp ann.
    (PrettyRep rep, Pretty inp) =>
    SubExp -> [inp] -> ScatterSpec VName -> Lambda rep -> Doc ann
    ppScatter SubExp
    w [VName]
    arrs ScatterSpec VName
    dests Lambda rep
    lam
      pretty (Hist SubExp
    w [VName]
    arrs [HistOp rep]
    ops Lambda rep
    bucket_fun) =
        SubExp -> [VName] -> [HistOp rep] -> Lambda rep -> Doc ann
    forall rep inp ann.
    (PrettyRep rep, Pretty inp) =>
    SubExp -> [inp] -> [HistOp rep] -> Lambda rep -> Doc ann
    ppHist SubExp
    w [VName]
    arrs [HistOp rep]
    ops Lambda rep
    bucket_fun
      pretty (Screma SubExp
    w [VName]
    arrs (ScremaForm Lambda rep
    map_lam [Scan rep]
    scans [Reduce rep]
    reds))
        | [Scan rep] -> Bool
    forall a. [a] -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null [Scan rep]
    scans,
          [Reduce rep] -> Bool
    forall a. [a] -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null [Reduce rep]
    reds =
            Doc ann
    "map"
              Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> (Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    parens (Doc ann -> Doc ann) -> (Doc ann -> Doc ann) -> Doc ann -> Doc ann
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    align)
                ( SubExp -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. SubExp -> Doc ann
    pretty SubExp
    w
                    Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    comma Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a
    </> [Doc ann] -> Doc ann
    forall a. [Doc a] -> Doc a
    ppTuple' ((VName -> Doc ann) -> [VName] -> [Doc ann]
    forall a b. (a -> b) -> [a] -> [b]
    map VName -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. VName -> Doc ann
    pretty [VName]
    arrs)
                    Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    comma Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a
    </> Lambda rep -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. Lambda rep -> Doc ann
    pretty Lambda rep
    map_lam
                )
        | [Scan rep] -> Bool
    forall a. [a] -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null [Scan rep]
    scans =
            Doc ann
    "redomap"
              Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> (Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    parens (Doc ann -> Doc ann) -> (Doc ann -> Doc ann) -> Doc ann -> Doc ann
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    align)
                ( SubExp -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. SubExp -> Doc ann
    pretty SubExp
    w
                    Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    comma Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a
    </> [Doc ann] -> Doc ann
    forall a. [Doc a] -> Doc a
    ppTuple' ((VName -> Doc ann) -> [VName] -> [Doc ann]
    forall a b. (a -> b) -> [a] -> [b]
    map VName -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. VName -> Doc ann
    pretty [VName]
    arrs)
                    Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    comma Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a
    </> Lambda rep -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. Lambda rep -> Doc ann
    pretty Lambda rep
    map_lam
                    Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    comma
                      Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a
    </> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    PP.braces ([Doc ann] -> Doc ann
    forall a. Monoid a => [a] -> a
    mconcat ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
    forall a b. (a -> b) -> a -> b
    $ Doc ann -> [Doc ann] -> [Doc ann]
    forall a. a -> [a] -> [a]
    intersperse (Doc ann
    forall ann. Doc ann
    comma Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    PP.line) ([Doc ann] -> [Doc ann]) -> [Doc ann] -> [Doc ann]
    forall a b. (a -> b) -> a -> b
    $ (Reduce rep -> Doc ann) -> [Reduce rep] -> [Doc ann]
    forall a b. (a -> b) -> [a] -> [b]
    map Reduce rep -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. Reduce rep -> Doc ann
    pretty [Reduce rep]
    reds)
                )
        | [Reduce rep] -> Bool
    forall a. [a] -> Bool
    forall (t :: * -> *) a. Foldable t => t a -> Bool
    null [Reduce rep]
    reds =
            Doc ann
    "scanomap"
              Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> (Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    parens (Doc ann -> Doc ann) -> (Doc ann -> Doc ann) -> Doc ann -> Doc ann
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    align)
                ( SubExp -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. SubExp -> Doc ann
    pretty SubExp
    w
                    Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    comma Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a
    </> [Doc ann] -> Doc ann
    forall a. [Doc a] -> Doc a
    ppTuple' ((VName -> Doc ann) -> [VName] -> [Doc ann]
    forall a b. (a -> b) -> [a] -> [b]
    map VName -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. VName -> Doc ann
    pretty [VName]
    arrs)
                    Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    comma Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a
    </> Lambda rep -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. Lambda rep -> Doc ann
    pretty Lambda rep
    map_lam
                    Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    comma
                      Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a
    </> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    PP.braces
                        ([Doc ann] -> Doc ann
    forall a. Monoid a => [a] -> a
    mconcat ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
    forall a b. (a -> b) -> a -> b
    $ Doc ann -> [Doc ann] -> [Doc ann]
    forall a. a -> [a] -> [a]
    intersperse (Doc ann
    forall ann. Doc ann
    comma Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    PP.line) ([Doc ann] -> [Doc ann]) -> [Doc ann] -> [Doc ann]
    forall a b. (a -> b) -> a -> b
    $ (Scan rep -> Doc ann) -> [Scan rep] -> [Doc ann]
    forall a b. (a -> b) -> [a] -> [b]
    map Scan rep -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. Scan rep -> Doc ann
    pretty [Scan rep]
    scans)
                )
      pretty (Screma SubExp
    w [VName]
    arrs ScremaForm rep
    form) = SubExp -> [VName] -> ScremaForm rep -> Doc ann
    forall rep inp ann.
    (PrettyRep rep, Pretty inp) =>
    SubExp -> [inp] -> ScremaForm rep -> Doc ann
    ppScrema SubExp
    w [VName]
    arrs ScremaForm rep
    form

    -- | Prettyprint the given Screma.
    ppScrema ::
      (PrettyRep rep, Pretty inp) => SubExp -> [inp] -> ScremaForm rep -> Doc ann
    ppScrema :: forall rep inp ann.
    (PrettyRep rep, Pretty inp) =>
    SubExp -> [inp] -> ScremaForm rep -> Doc ann
    ppScrema SubExp
    w [inp]
    arrs (ScremaForm Lambda rep
    map_lam [Scan rep]
    scans [Reduce rep]
    reds) =
      Doc ann
    "screma"
        Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> (Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    parens (Doc ann -> Doc ann) -> (Doc ann -> Doc ann) -> Doc ann -> Doc ann
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    align)
          ( SubExp -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. SubExp -> Doc ann
    pretty SubExp
    w
              Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    comma Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a
    </> [Doc ann] -> Doc ann
    forall a. [Doc a] -> Doc a
    ppTuple' ((inp -> Doc ann) -> [inp] -> [Doc ann]
    forall a b. (a -> b) -> [a] -> [b]
    map inp -> Doc ann
    forall ann. inp -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    pretty [inp]
    arrs)
              Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    comma Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a
    </> Lambda rep -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. Lambda rep -> Doc ann
    pretty Lambda rep
    map_lam
              Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    comma
                Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a
    </> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    PP.braces ([Doc ann] -> Doc ann
    forall a. Monoid a => [a] -> a
    mconcat ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
    forall a b. (a -> b) -> a -> b
    $ Doc ann -> [Doc ann] -> [Doc ann]
    forall a. a -> [a] -> [a]
    intersperse (Doc ann
    forall ann. Doc ann
    comma Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    PP.line) ([Doc ann] -> [Doc ann]) -> [Doc ann] -> [Doc ann]
    forall a b. (a -> b) -> a -> b
    $ (Scan rep -> Doc ann) -> [Scan rep] -> [Doc ann]
    forall a b. (a -> b) -> [a] -> [b]
    map Scan rep -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. Scan rep -> Doc ann
    pretty [Scan rep]
    scans)
              Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    comma
                Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a
    </> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    PP.braces ([Doc ann] -> Doc ann
    forall a. Monoid a => [a] -> a
    mconcat ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
    forall a b. (a -> b) -> a -> b
    $ Doc ann -> [Doc ann] -> [Doc ann]
    forall a. a -> [a] -> [a]
    intersperse (Doc ann
    forall ann. Doc ann
    comma Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    PP.line) ([Doc ann] -> [Doc ann]) -> [Doc ann] -> [Doc ann]
    forall a b. (a -> b) -> a -> b
    $ (Reduce rep -> Doc ann) -> [Reduce rep] -> [Doc ann]
    forall a b. (a -> b) -> [a] -> [b]
    map Reduce rep -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. Reduce rep -> Doc ann
    pretty [Reduce rep]
    reds)
          )

    -- | Prettyprint the given Stream.
    ppStream ::
      (PrettyRep rep, Pretty inp) => SubExp -> [inp] -> [SubExp] -> Lambda rep -> Doc ann
    ppStream :: forall rep inp ann.
    (PrettyRep rep, Pretty inp) =>
    SubExp -> [inp] -> [SubExp] -> Lambda rep -> Doc ann
    ppStream SubExp
    size [inp]
    arrs [SubExp]
    acc Lambda rep
    lam =
      Doc ann
    "streamSeq"
        Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> (Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    parens (Doc ann -> Doc ann) -> (Doc ann -> Doc ann) -> Doc ann -> Doc ann
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    align)
          ( SubExp -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. SubExp -> Doc ann
    pretty SubExp
    size
              Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    comma
                Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a
    </> [Doc ann] -> Doc ann
    forall a. [Doc a] -> Doc a
    ppTuple' ((inp -> Doc ann) -> [inp] -> [Doc ann]
    forall a b. (a -> b) -> [a] -> [b]
    map inp -> Doc ann
    forall ann. inp -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    pretty [inp]
    arrs)
              Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    comma
                Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a
    </> [Doc ann] -> Doc ann
    forall a. [Doc a] -> Doc a
    ppTuple' ((SubExp -> Doc ann) -> [SubExp] -> [Doc ann]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExp -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. SubExp -> Doc ann
    pretty [SubExp]
    acc)
              Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    comma
                Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a
    </> Lambda rep -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. Lambda rep -> Doc ann
    pretty Lambda rep
    lam
          )

    -- | Prettyprint the given Scatter.
    ppScatter ::
      (PrettyRep rep, Pretty inp) => SubExp -> [inp] -> [(Shape, Int, VName)] -> Lambda rep -> Doc ann
    ppScatter :: forall rep inp ann.
    (PrettyRep rep, Pretty inp) =>
    SubExp -> [inp] -> ScatterSpec VName -> Lambda rep -> Doc ann
    ppScatter SubExp
    w [inp]
    arrs ScatterSpec VName
    dests Lambda rep
    lam =
      Doc ann
    "scatter"
        Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> (Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    parens (Doc ann -> Doc ann) -> (Doc ann -> Doc ann) -> Doc ann -> Doc ann
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
    Category cat =>
    cat b c -> cat a b -> cat a c
    . Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    align)
          ( SubExp -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. SubExp -> Doc ann
    pretty SubExp
    w
              Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    comma
                Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a
    </> [Doc ann] -> Doc ann
    forall a. [Doc a] -> Doc a
    ppTuple' ((inp -> Doc ann) -> [inp] -> [Doc ann]
    forall a b. (a -> b) -> [a] -> [b]
    map inp -> Doc ann
    forall ann. inp -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    pretty [inp]
    arrs)
              Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    comma
                Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a
    </> [Doc ann] -> Doc ann
    forall a. [Doc a] -> Doc a
    commasep (((Shape, Int, VName) -> Doc ann) -> ScatterSpec VName -> [Doc ann]
    forall a b. (a -> b) -> [a] -> [b]
    map (Shape, Int, VName) -> Doc ann
    forall ann. (Shape, Int, VName) -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    pretty ScatterSpec VName
    dests)
              Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    comma
                Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a
    </> Lambda rep -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. Lambda rep -> Doc ann
    pretty Lambda rep
    lam
          )

    instance (PrettyRep rep) => Pretty (Scan rep) where
      pretty :: forall ann. Scan rep -> Doc ann
    pretty (Scan Lambda rep
    scan_lam [SubExp]
    scan_nes) =
        Lambda rep -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. Lambda rep -> Doc ann
    pretty Lambda rep
    scan_lam Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    comma Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a
    </> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    PP.braces ([Doc ann] -> Doc ann
    forall a. [Doc a] -> Doc a
    commasep ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
    forall a b. (a -> b) -> a -> b
    $ (SubExp -> Doc ann) -> [SubExp] -> [Doc ann]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExp -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. SubExp -> Doc ann
    pretty [SubExp]
    scan_nes)

    ppComm :: Commutativity -> Doc ann
    ppComm :: forall ann. Commutativity -> Doc ann
    ppComm Commutativity
    Noncommutative = Doc ann
    forall a. Monoid a => a
    mempty
    ppComm Commutativity
    Commutative = Doc ann
    "commutative "

    instance (PrettyRep rep) => Pretty (Reduce rep) where
      pretty :: forall ann. Reduce rep -> Doc ann
    pretty (Reduce Commutativity
    comm Lambda rep
    red_lam [SubExp]
    red_nes) =
        Commutativity -> Doc ann
    forall ann. Commutativity -> Doc ann
    ppComm Commutativity
    comm
          Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Lambda rep -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. Lambda rep -> Doc ann
    pretty Lambda rep
    red_lam
          Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    comma
            Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a
    </> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    PP.braces ([Doc ann] -> Doc ann
    forall a. [Doc a] -> Doc a
    commasep ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
    forall a b. (a -> b) -> a -> b
    $ (SubExp -> Doc ann) -> [SubExp] -> [Doc ann]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExp -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. SubExp -> Doc ann
    pretty [SubExp]
    red_nes)

    -- | Prettyprint the given histogram operation.
    ppHist ::
      (PrettyRep rep, Pretty inp) =>
      SubExp ->
      [inp] ->
      [HistOp rep] ->
      Lambda rep ->
      Doc ann
    ppHist :: forall rep inp ann.
    (PrettyRep rep, Pretty inp) =>
    SubExp -> [inp] -> [HistOp rep] -> Lambda rep -> Doc ann
    ppHist SubExp
    w [inp]
    arrs [HistOp rep]
    ops Lambda rep
    bucket_fun =
      Doc ann
    "hist"
        Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    parens
          ( SubExp -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. SubExp -> Doc ann
    pretty SubExp
    w
              Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    comma
                Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a
    </> [Doc ann] -> Doc ann
    forall a. [Doc a] -> Doc a
    ppTuple' ((inp -> Doc ann) -> [inp] -> [Doc ann]
    forall a b. (a -> b) -> [a] -> [b]
    map inp -> Doc ann
    forall ann. inp -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    pretty [inp]
    arrs)
              Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    comma
                Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a
    </> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    PP.braces ([Doc ann] -> Doc ann
    forall a. Monoid a => [a] -> a
    mconcat ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
    forall a b. (a -> b) -> a -> b
    $ Doc ann -> [Doc ann] -> [Doc ann]
    forall a. a -> [a] -> [a]
    intersperse (Doc ann
    forall ann. Doc ann
    comma Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    PP.line) ([Doc ann] -> [Doc ann]) -> [Doc ann] -> [Doc ann]
    forall a b. (a -> b) -> a -> b
    $ (HistOp rep -> Doc ann) -> [HistOp rep] -> [Doc ann]
    forall a b. (a -> b) -> [a] -> [b]
    map HistOp rep -> Doc ann
    forall {rep} {ann}. PrettyRep rep => HistOp rep -> Doc ann
    ppOp [HistOp rep]
    ops)
              Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    comma
                Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a
    </> Lambda rep -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. Lambda rep -> Doc ann
    pretty Lambda rep
    bucket_fun
          )
      where
        ppOp :: HistOp rep -> Doc ann
    ppOp (HistOp Shape
    dest_w SubExp
    rf [VName]
    dests [SubExp]
    nes Lambda rep
    op) =
          Shape -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. Shape -> Doc ann
    pretty Shape
    dest_w
            Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    comma
              Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a
    <+> SubExp -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. SubExp -> Doc ann
    pretty SubExp
    rf
            Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    comma
              Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a
    <+> Doc ann -> Doc ann
    forall ann. Doc ann -> Doc ann
    PP.braces ([Doc ann] -> Doc ann
    forall a. [Doc a] -> Doc a
    commasep ([Doc ann] -> Doc ann) -> [Doc ann] -> Doc ann
    forall a b. (a -> b) -> a -> b
    $ (VName -> Doc ann) -> [VName] -> [Doc ann]
    forall a b. (a -> b) -> [a] -> [b]
    map VName -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. VName -> Doc ann
    pretty [VName]
    dests)
            Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    comma
              Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a
    </> [Doc ann] -> Doc ann
    forall a. [Doc a] -> Doc a
    ppTuple' ((SubExp -> Doc ann) -> [SubExp] -> [Doc ann]
    forall a b. (a -> b) -> [a] -> [b]
    map SubExp -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. SubExp -> Doc ann
    pretty [SubExp]
    nes)
            Doc ann -> Doc ann -> Doc ann
    forall a. Semigroup a => a -> a -> a
    <> Doc ann
    forall ann. Doc ann
    comma
              Doc ann -> Doc ann -> Doc ann
    forall a. Doc a -> Doc a -> Doc a
    </> Lambda rep -> Doc ann
    forall a ann. Pretty a => a -> Doc ann
    forall ann. Lambda rep -> Doc ann
    pretty Lambda rep
    op