{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE ConstraintKinds     #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeFamilies        #-}
{-# LANGUAGE TypeOperators       #-}
{-# LANGUAGE ViewPatterns        #-}
-- |
-- Module      : Data.Array.Accelerate.Language
-- Copyright   : [2008..2020] The Accelerate Team
-- License     : BSD3
--
-- Maintainer  : Trevor L. McDonell <trevor.mcdonell@gmail.com>
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--
-- We use the dictionary view of overloaded operations (such as arithmetic and
-- bit manipulation) to reify such expressions.  With non-overloaded
-- operations (such as, the logical connectives) and partially overloaded
-- operations (such as comparisons), we use the standard operator names with a
-- \'*\' attached.  We keep the standard alphanumeric names as they can be
-- easily qualified.
--

module Data.Array.Accelerate.Language (

  -- * Array construction
  use, unit, replicate, generate,

  -- * Shape manipulation
  reshape,

  -- * Extraction of sub-arrays
  slice,

  -- * Map-like functions
  map, zipWith,

  -- -- * Sequence collection
  -- collect,

  -- -- * Sequence producers
  -- streamIn, toSeq,

  -- -- * Sequence transducers
  -- mapSeq, zipWithSeq, scanSeq,

  -- -- * Sequence consumers
  -- foldSeq, foldSeqFlatten,

  -- * Reductions
  fold, fold1, foldSeg', fold1Seg',

  -- * Scan functions
  scanl, scanl', scanl1, scanr, scanr', scanr1,

  -- * Permutations
  permute, backpermute,

  -- * Stencil operations
  stencil, stencil2,

  -- ** Stencil specification
  Boundary, Stencil,
  clamp, mirror, wrap, function,


  -- ** Common stencil types
  Stencil3, Stencil5, Stencil7, Stencil9,
  Stencil3x3, Stencil5x3, Stencil3x5, Stencil5x5,
  Stencil3x3x3, Stencil5x3x3, Stencil3x5x3, Stencil3x3x5, Stencil5x5x3, Stencil5x3x5,
  Stencil3x5x5, Stencil5x5x5,

  -- * Foreign functions
  foreignAcc,
  foreignExp,

  -- * Pipelining
  (>->),

  -- * Index construction and destruction
  indexHead, indexTail, toIndex, fromIndex,
  intersect, union,

  -- * Flow-control
  acond, awhile,
  cond,  while,

  -- * Array operations with a scalar result
  (!), (!!), shape, size, shapeSize,

  -- * Numeric functions
  subtract, even, odd, gcd, lcm, (^), (^^),

  -- * Conversions
  ord, chr, boolToInt, bitcast,

) where

import Data.Array.Accelerate.AST                                    ( PrimFun(..) )
import Data.Array.Accelerate.Pattern
import Data.Array.Accelerate.Representation.Array                   ( ArrayR(..) )
import Data.Array.Accelerate.Representation.Shape                   ( ShapeR(..) )
import Data.Array.Accelerate.Representation.Type
import Data.Array.Accelerate.Smart                                  hiding ( arraysR )
import Data.Array.Accelerate.Sugar.Array                            ( Arrays(..), Array, Scalar, Segments, arrayR )
import Data.Array.Accelerate.Sugar.Elt
import Data.Array.Accelerate.Sugar.Foreign
import Data.Array.Accelerate.Sugar.Shape                            ( Shape(..), Slice(..), (:.) )
import Data.Array.Accelerate.Type
import qualified Data.Array.Accelerate.Representation.Array         as R

import Data.Array.Accelerate.Classes.Eq
import Data.Array.Accelerate.Classes.Fractional
import Data.Array.Accelerate.Classes.Integral
import Data.Array.Accelerate.Classes.Num
import Data.Array.Accelerate.Classes.Ord

import Prelude                                                      ( ($), (.), Maybe(..), Char )


-- $setup
-- >>> :seti -XFlexibleContexts
-- >>> :seti -XScopedTypeVariables
-- >>> :seti -XTypeOperators
-- >>> :seti -XViewPatterns
-- >>> import Data.Array.Accelerate
-- >>> import Data.Array.Accelerate.Data.Maybe
-- >>> import Data.Array.Accelerate.Interpreter
-- >>> :{
--   let runExp :: Elt e => Exp e -> e
--       runExp e = indexArray (run (unit e)) Z
-- :}

-- Array introduction
-- ------------------

-- | Make an array from vanilla Haskell available for processing within embedded
-- Accelerate computations.
--
-- Depending upon which backend is used to eventually execute array
-- computations, 'use' may entail data transfer (e.g. to a GPU).
--
-- 'use' is overloaded so that it can accept tuples of 'Arrays':
--
-- >>> let vec = fromList (Z:.10) [0..] :: Vector Int
-- >>> vec
-- Vector (Z :. 10) [0,1,2,3,4,5,6,7,8,9]
--
-- >>> let mat = fromList (Z:.5:.10) [0..] :: Matrix Int
-- >>> mat
-- Matrix (Z :. 5 :. 10)
--   [  0,  1,  2,  3,  4,  5,  6,  7,  8,  9,
--     10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
--     20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
--     30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
--     40, 41, 42, 43, 44, 45, 46, 47, 48, 49]
--
-- >>> let vec' = use vec         :: Acc (Vector Int)
-- >>> let mat' = use mat         :: Acc (Matrix Int)
-- >>> let tup  = use (vec, mat)  :: Acc (Vector Int, Matrix Int)
--
use :: forall arrays. Arrays arrays => arrays -> Acc arrays
use :: arrays -> Acc arrays
use = SmartAcc (ArraysR arrays) -> Acc arrays
forall a. SmartAcc (ArraysR a) -> Acc a
Acc (SmartAcc (ArraysR arrays) -> Acc arrays)
-> (arrays -> SmartAcc (ArraysR arrays)) -> arrays -> Acc arrays
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ArraysR (ArraysR arrays)
-> ArraysR arrays -> SmartAcc (ArraysR arrays)
forall a. ArraysR a -> a -> SmartAcc a
use' (Arrays arrays => ArraysR (ArraysR arrays)
forall a. Arrays a => ArraysR (ArraysR a)
arraysR @arrays) (ArraysR arrays -> SmartAcc (ArraysR arrays))
-> (arrays -> ArraysR arrays)
-> arrays
-> SmartAcc (ArraysR arrays)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. arrays -> ArraysR arrays
forall a. Arrays a => a -> ArraysR a
fromArr
  where
    use' :: R.ArraysR a -> a -> SmartAcc a
    use' :: ArraysR a -> a -> SmartAcc a
use' ArraysR a
TupRunit                   ()       = PreSmartAcc SmartAcc SmartExp () -> SmartAcc ()
forall a. PreSmartAcc SmartAcc SmartExp a -> SmartAcc a
SmartAcc (PreSmartAcc SmartAcc SmartExp () -> SmartAcc ())
-> PreSmartAcc SmartAcc SmartExp () -> SmartAcc ()
forall a b. (a -> b) -> a -> b
$ PreSmartAcc SmartAcc SmartExp ()
forall (acc :: * -> *) (exp :: * -> *). PreSmartAcc acc exp ()
Anil
    use' (TupRsingle repr :: ArrayR a
repr@ArrayR{}) a
a        = PreSmartAcc SmartAcc SmartExp (Array sh e) -> SmartAcc (Array sh e)
forall a. PreSmartAcc SmartAcc SmartExp a -> SmartAcc a
SmartAcc (PreSmartAcc SmartAcc SmartExp (Array sh e)
 -> SmartAcc (Array sh e))
-> PreSmartAcc SmartAcc SmartExp (Array sh e)
-> SmartAcc (Array sh e)
forall a b. (a -> b) -> a -> b
$ ArrayR (Array sh e)
-> Array sh e -> PreSmartAcc SmartAcc SmartExp (Array sh e)
forall sh e (acc :: * -> *) (exp :: * -> *).
ArrayR (Array sh e)
-> Array sh e -> PreSmartAcc acc exp (Array sh e)
Use ArrayR a
ArrayR (Array sh e)
repr a
Array sh e
a
    use' (TupRpair TupR ArrayR a
r1 TupR ArrayR b
r2)           (a1, a2) = PreSmartAcc SmartAcc SmartExp (a, b) -> SmartAcc (a, b)
forall a. PreSmartAcc SmartAcc SmartExp a -> SmartAcc a
SmartAcc (PreSmartAcc SmartAcc SmartExp (a, b) -> SmartAcc (a, b))
-> PreSmartAcc SmartAcc SmartExp (a, b) -> SmartAcc (a, b)
forall a b. (a -> b) -> a -> b
$ TupR ArrayR a -> a -> SmartAcc a
forall a. ArraysR a -> a -> SmartAcc a
use' TupR ArrayR a
r1 a
a1 SmartAcc a -> SmartAcc b -> PreSmartAcc SmartAcc SmartExp (a, b)
forall (acc :: * -> *) arrs1 arrs2 (exp :: * -> *).
acc arrs1 -> acc arrs2 -> PreSmartAcc acc exp (arrs1, arrs2)
`Apair` TupR ArrayR b -> b -> SmartAcc b
forall a. ArraysR a -> a -> SmartAcc a
use' TupR ArrayR b
r2 b
a2

-- | Construct a singleton (one element) array from a scalar value (or tuple of
-- scalar values).
--
unit :: forall e. Elt e => Exp e -> Acc (Scalar e)
unit :: Exp e -> Acc (Scalar e)
unit (Exp SmartExp (EltR e)
e) = SmartAcc (ArraysR (Scalar e)) -> Acc (Scalar e)
forall a. SmartAcc (ArraysR a) -> Acc a
Acc (SmartAcc (ArraysR (Scalar e)) -> Acc (Scalar e))
-> SmartAcc (ArraysR (Scalar e)) -> Acc (Scalar e)
forall a b. (a -> b) -> a -> b
$ PreSmartAcc SmartAcc SmartExp (Scalar (EltR e))
-> SmartAcc (Scalar (EltR e))
forall a. PreSmartAcc SmartAcc SmartExp a -> SmartAcc a
SmartAcc (PreSmartAcc SmartAcc SmartExp (Scalar (EltR e))
 -> SmartAcc (Scalar (EltR e)))
-> PreSmartAcc SmartAcc SmartExp (Scalar (EltR e))
-> SmartAcc (Scalar (EltR e))
forall a b. (a -> b) -> a -> b
$ TypeR (EltR e)
-> SmartExp (EltR e)
-> PreSmartAcc SmartAcc SmartExp (Scalar (EltR e))
forall e (exp :: * -> *) (acc :: * -> *).
TypeR e -> exp e -> PreSmartAcc acc exp (Scalar e)
Unit (Elt e => TypeR (EltR e)
forall a. Elt a => TypeR (EltR a)
eltR @e) SmartExp (EltR e)
e

-- | Replicate an array across one or more dimensions as specified by the
-- /generalised/ array index provided as the first argument.
--
-- For example, given the following vector:
--
-- >>> let vec = fromList (Z:.10) [0..] :: Vector Int
-- >>> vec
-- Vector (Z :. 10) [0,1,2,3,4,5,6,7,8,9]
--
-- ...we can replicate these elements to form a two-dimensional array either by
-- replicating those elements as new rows:
--
-- >>> run $ replicate (constant (Z :. (4::Int) :. All)) (use vec)
-- Matrix (Z :. 4 :. 10)
--   [ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
--     0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
--     0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
--     0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
--
-- ...or as columns:
--
-- >>> run $ replicate (lift (Z :. All :. (4::Int))) (use vec)
-- Matrix (Z :. 10 :. 4)
--   [ 0, 0, 0, 0,
--     1, 1, 1, 1,
--     2, 2, 2, 2,
--     3, 3, 3, 3,
--     4, 4, 4, 4,
--     5, 5, 5, 5,
--     6, 6, 6, 6,
--     7, 7, 7, 7,
--     8, 8, 8, 8,
--     9, 9, 9, 9]
--
-- Replication along more than one dimension is also possible. Here we replicate
-- twice across the first dimension and three times across the third dimension:
--
-- >>> run $ replicate (constant (Z :. (2::Int) :. All :. (3::Int))) (use vec)
-- Array (Z :. 2 :. 10 :. 3) [0,0,0,1,1,1,2,2,2,3,3,3,4,4,4,5,5,5,6,6,6,7,7,7,8,8,8,9,9,9,0,0,0,1,1,1,2,2,2,3,3,3,4,4,4,5,5,5,6,6,6,7,7,7,8,8,8,9,9,9]
--
-- The marker 'Any' can be used in the slice specification to match against some
-- arbitrary dimension. For example, here 'Any' matches against whatever shape
-- type variable @sh@ takes.
--
-- >>> :{
--   let rep0 :: (Shape sh, Elt e) => Exp Int -> Acc (Array sh e) -> Acc (Array (sh :. Int) e)
--       rep0 n a = replicate (lift (Any :. n)) a
-- :}
--
-- >>> let x = unit 42 :: Acc (Scalar Int)
-- >>> run $ rep0 10 x
-- Vector (Z :. 10) [42,42,42,42,42,42,42,42,42,42]
--
-- >>> run $ rep0 5 (use vec)
-- Matrix (Z :. 10 :. 5)
--   [ 0, 0, 0, 0, 0,
--     1, 1, 1, 1, 1,
--     2, 2, 2, 2, 2,
--     3, 3, 3, 3, 3,
--     4, 4, 4, 4, 4,
--     5, 5, 5, 5, 5,
--     6, 6, 6, 6, 6,
--     7, 7, 7, 7, 7,
--     8, 8, 8, 8, 8,
--     9, 9, 9, 9, 9]
--
-- Of course, 'Any' and 'All' can be used together.
--
-- >>> :{
--   let rep1 :: (Shape sh, Elt e) => Exp Int -> Acc (Array (sh :. Int) e) -> Acc (Array (sh :. Int :. Int) e)
--       rep1 n a = replicate (lift (Any :. n :. All)) a
-- :}
--
-- >>> run $ rep1 5 (use vec)
-- Matrix (Z :. 5 :. 10)
--   [ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
--     0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
--     0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
--     0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
--     0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
--
replicate
    :: forall slix e.
       (Slice slix, Elt e)
    => Exp slix
    -> Acc (Array (SliceShape slix) e)
    -> Acc (Array (FullShape  slix) e)
replicate :: Exp slix
-> Acc (Array (SliceShape slix) e)
-> Acc (Array (FullShape slix) e)
replicate = SmartAcc (Array (EltR (FullShape slix)) (EltR e))
-> Acc (Array (FullShape slix) e)
forall a. SmartAcc (ArraysR a) -> Acc a
Acc (SmartAcc (Array (EltR (FullShape slix)) (EltR e))
 -> Acc (Array (FullShape slix) e))
-> (Exp slix
    -> Acc (Array (SliceShape slix) e)
    -> SmartAcc (Array (EltR (FullShape slix)) (EltR e)))
-> Exp slix
-> Acc (Array (SliceShape slix) e)
-> Acc (Array (FullShape slix) e)
forall b a c d. (b -> a) -> (c -> d -> b) -> c -> d -> a
$$ FromApplyAcc
  (Exp slix
   -> Acc (Array (SliceShape slix) e)
   -> SmartAcc (Array (EltR (FullShape slix)) (EltR e)))
-> Exp slix
-> Acc (Array (SliceShape slix) e)
-> SmartAcc (Array (EltR (FullShape slix)) (EltR e))
forall a. ApplyAcc a => FromApplyAcc a -> a
applyAcc (SliceIndex
  (EltR slix)
  (EltR (SliceShape slix))
  (EltR (CoSliceShape slix))
  (EltR (FullShape slix))
-> SmartExp (EltR slix)
-> SmartAcc (Array (EltR (SliceShape slix)) (EltR e))
-> PreSmartAcc
     SmartAcc SmartExp (Array (EltR (FullShape slix)) (EltR e))
forall slix sl co sh (exp :: * -> *) (acc :: * -> *) e.
SliceIndex slix sl co sh
-> exp slix -> acc (Array sl e) -> PreSmartAcc acc exp (Array sh e)
Replicate (SliceIndex
   (EltR slix)
   (EltR (SliceShape slix))
   (EltR (CoSliceShape slix))
   (EltR (FullShape slix))
 -> SmartExp (EltR slix)
 -> SmartAcc (Array (EltR (SliceShape slix)) (EltR e))
 -> PreSmartAcc
      SmartAcc SmartExp (Array (EltR (FullShape slix)) (EltR e)))
-> SliceIndex
     (EltR slix)
     (EltR (SliceShape slix))
     (EltR (CoSliceShape slix))
     (EltR (FullShape slix))
-> SmartExp (EltR slix)
-> SmartAcc (Array (EltR (SliceShape slix)) (EltR e))
-> PreSmartAcc
     SmartAcc SmartExp (Array (EltR (FullShape slix)) (EltR e))
forall a b. (a -> b) -> a -> b
$ Slice slix =>
SliceIndex
  (EltR slix)
  (EltR (SliceShape slix))
  (EltR (CoSliceShape slix))
  (EltR (FullShape slix))
forall sl.
Slice sl =>
SliceIndex
  (EltR sl)
  (EltR (SliceShape sl))
  (EltR (CoSliceShape sl))
  (EltR (FullShape sl))
sliceIndex @slix)

-- | Construct a new array by applying a function to each index.
--
-- For example, the following will generate a one-dimensional array
-- (`Vector`) of three floating point numbers:
--
-- >>> run $ generate (I1 3) (\_ -> 1.2) :: Vector Float
-- Vector (Z :. 3) [1.2,1.2,1.2]
--
-- Or equivalently:
--
-- >>> run $ fill (constant (Z :. 3)) 1.2 :: Vector Float
-- Vector (Z :. 3) [1.2,1.2,1.2]
--
-- The following will create a vector with the elements @[1..10]@:
--
-- >>> run $ generate (I1 10) (\(I1 i) -> i + 1) :: Vector Int
-- Vector (Z :. 10) [1,2,3,4,5,6,7,8,9,10]
--
-- [/NOTE:/]
--
-- Using 'generate', it is possible to introduce nested data parallelism, which
-- will cause the program to fail.
--
-- If the index given by the scalar function is then used to dispatch further
-- parallel work, whose result is returned into 'Exp' terms by array indexing
-- operations such as ('!') or 'Data.Array.Accelerate.Prelude.the', the program
-- will fail with the error:
-- @.\/Data\/Array\/Accelerate\/Trafo\/Sharing.hs:447 (convertSharingExp): inconsistent valuation \@ shared \'Exp\' tree ...@.
--
generate
    :: forall sh a.
       (Shape sh, Elt a)
    => Exp sh
    -> (Exp sh -> Exp a)
    -> Acc (Array sh a)
generate :: Exp sh -> (Exp sh -> Exp a) -> Acc (Array sh a)
generate = SmartAcc (Array (EltR sh) (EltR a)) -> Acc (Array sh a)
forall a. SmartAcc (ArraysR a) -> Acc a
Acc (SmartAcc (Array (EltR sh) (EltR a)) -> Acc (Array sh a))
-> (Exp sh
    -> (Exp sh -> Exp a) -> SmartAcc (Array (EltR sh) (EltR a)))
-> Exp sh
-> (Exp sh -> Exp a)
-> Acc (Array sh a)
forall b a c d. (b -> a) -> (c -> d -> b) -> c -> d -> a
$$ FromApplyAcc
  (Exp sh
   -> (Exp sh -> Exp a) -> SmartAcc (Array (EltR sh) (EltR a)))
-> Exp sh
-> (Exp sh -> Exp a)
-> SmartAcc (Array (EltR sh) (EltR a))
forall a. ApplyAcc a => FromApplyAcc a -> a
applyAcc (ArrayR (Array (EltR sh) (EltR a))
-> SmartExp (EltR sh)
-> (SmartExp (EltR sh) -> SmartExp (EltR a))
-> PreSmartAcc SmartAcc SmartExp (Array (EltR sh) (EltR a))
forall sh e (exp :: * -> *) (acc :: * -> *).
ArrayR (Array sh e)
-> exp sh
-> (SmartExp sh -> exp e)
-> PreSmartAcc acc exp (Array sh e)
Generate (ArrayR (Array (EltR sh) (EltR a))
 -> SmartExp (EltR sh)
 -> (SmartExp (EltR sh) -> SmartExp (EltR a))
 -> PreSmartAcc SmartAcc SmartExp (Array (EltR sh) (EltR a)))
-> ArrayR (Array (EltR sh) (EltR a))
-> SmartExp (EltR sh)
-> (SmartExp (EltR sh) -> SmartExp (EltR a))
-> PreSmartAcc SmartAcc SmartExp (Array (EltR sh) (EltR a))
forall a b. (a -> b) -> a -> b
$ (Shape sh, Elt a) => ArrayR (Array (EltR sh) (EltR a))
forall sh e. (Shape sh, Elt e) => ArrayR (Array (EltR sh) (EltR e))
arrayR @sh @a)

-- Shape manipulation
-- ------------------

-- | Change the shape of an array without altering its contents. The 'size' of
-- the source and result arrays must be identical.
--
-- > precondition: shapeSize sh == shapeSize sh'
--
-- If the argument array is manifest in memory, 'reshape' is a no-op. If the
-- argument is to be fused into a subsequent operation, 'reshape' corresponds to
-- an index transformation in the fused code.
--
reshape
    :: forall sh sh' e.
       (Shape sh, Shape sh', Elt e)
    => Exp sh
    -> Acc (Array sh' e)
    -> Acc (Array sh e)
reshape :: Exp sh -> Acc (Array sh' e) -> Acc (Array sh e)
reshape = SmartAcc (Array (EltR sh) (EltR e)) -> Acc (Array sh e)
forall a. SmartAcc (ArraysR a) -> Acc a
Acc (SmartAcc (Array (EltR sh) (EltR e)) -> Acc (Array sh e))
-> (Exp sh
    -> Acc (Array sh' e) -> SmartAcc (Array (EltR sh) (EltR e)))
-> Exp sh
-> Acc (Array sh' e)
-> Acc (Array sh e)
forall b a c d. (b -> a) -> (c -> d -> b) -> c -> d -> a
$$ FromApplyAcc
  (Exp sh
   -> Acc (Array sh' e) -> SmartAcc (Array (EltR sh) (EltR e)))
-> Exp sh
-> Acc (Array sh' e)
-> SmartAcc (Array (EltR sh) (EltR e))
forall a. ApplyAcc a => FromApplyAcc a -> a
applyAcc (ShapeR (EltR sh)
-> SmartExp (EltR sh)
-> SmartAcc (Array (EltR sh') (EltR e))
-> PreSmartAcc SmartAcc SmartExp (Array (EltR sh) (EltR e))
forall sh (exp :: * -> *) (acc :: * -> *) sh' e.
ShapeR sh
-> exp sh -> acc (Array sh' e) -> PreSmartAcc acc exp (Array sh e)
Reshape (ShapeR (EltR sh)
 -> SmartExp (EltR sh)
 -> SmartAcc (Array (EltR sh') (EltR e))
 -> PreSmartAcc SmartAcc SmartExp (Array (EltR sh) (EltR e)))
-> ShapeR (EltR sh)
-> SmartExp (EltR sh)
-> SmartAcc (Array (EltR sh') (EltR e))
-> PreSmartAcc SmartAcc SmartExp (Array (EltR sh) (EltR e))
forall a b. (a -> b) -> a -> b
$ Shape sh => ShapeR (EltR sh)
forall sh. Shape sh => ShapeR (EltR sh)
shapeR @sh)

-- Extraction of sub-arrays
-- ------------------------

-- | Index an array with a /generalised/ array index, supplied as the second
-- argument. The result is a new array (possibly a singleton) containing the
-- selected dimensions ('All's) in their entirety.
--
-- 'slice' is the opposite of 'replicate', and can be used to /cut out/ entire
-- dimensions. For example, for the two dimensional array 'mat':
--
-- >>> let mat = fromList (Z:.5:.10) [0..] :: Matrix Int
-- >>> mat
-- Matrix (Z :. 5 :. 10)
--   [  0,  1,  2,  3,  4,  5,  6,  7,  8,  9,
--     10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
--     20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
--     30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
--     40, 41, 42, 43, 44, 45, 46, 47, 48, 49]
--
-- ...will can select a specific row to yield a one dimensional result by fixing
-- the row index (2) while allowing the column index to vary (via 'All'):
--
-- >>> run $ slice (use mat) (constant (Z :. (2::Int) :. All))
-- Vector (Z :. 10) [20,21,22,23,24,25,26,27,28,29]
--
-- A fully specified index (with no 'All's) returns a single element (zero
-- dimensional array).
--
-- >>> run $ slice (use mat) (constant (Z :. 4 :. 2 :: DIM2))
-- Scalar Z [42]
--
-- The marker 'Any' can be used in the slice specification to match against some
-- arbitrary (lower) dimension. Here 'Any' matches whatever shape type variable
-- @sh@ takes:
--
-- >>> :{
--   let
--       sl0 :: (Shape sh, Elt e) => Acc (Array (sh:.Int) e) -> Exp Int -> Acc (Array sh e)
--       sl0 a n = slice a (lift (Any :. n))
-- :}
--
-- >>> let vec = fromList (Z:.10) [0..] :: Vector Int
-- >>> run $ sl0 (use vec) 4
-- Scalar Z [4]
--
-- >>> run $ sl0 (use mat) 4
-- Vector (Z :. 5) [4,14,24,34,44]
--
-- Of course, 'Any' and 'All' can be used together.
--
-- >>> :{
--   let sl1 :: (Shape sh, Elt e) => Acc (Array (sh:.Int:.Int) e) -> Exp Int -> Acc (Array (sh:.Int) e)
--       sl1 a n = slice a (lift (Any :. n :. All))
-- :}
--
-- >>> run $ sl1 (use mat) 4
-- Vector (Z :. 10) [40,41,42,43,44,45,46,47,48,49]
--
-- >>> let cube = fromList (Z:.3:.4:.5) [0..] :: Array DIM3 Int
-- >>> cube
-- Array (Z :. 3 :. 4 :. 5) [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59]
--
-- >>> run $ sl1 (use cube) 2
-- Matrix (Z :. 3 :. 5)
--   [ 10, 11, 12, 13, 14,
--     30, 31, 32, 33, 34,
--     50, 51, 52, 53, 54]
--
slice :: forall slix e.
         (Slice slix, Elt e)
      => Acc (Array (FullShape slix) e)
      -> Exp slix
      -> Acc (Array (SliceShape slix) e)
slice :: Acc (Array (FullShape slix) e)
-> Exp slix -> Acc (Array (SliceShape slix) e)
slice = SmartAcc (Array (EltR (SliceShape slix)) (EltR e))
-> Acc (Array (SliceShape slix) e)
forall a. SmartAcc (ArraysR a) -> Acc a
Acc (SmartAcc (Array (EltR (SliceShape slix)) (EltR e))
 -> Acc (Array (SliceShape slix) e))
-> (Acc (Array (FullShape slix) e)
    -> Exp slix -> SmartAcc (Array (EltR (SliceShape slix)) (EltR e)))
-> Acc (Array (FullShape slix) e)
-> Exp slix
-> Acc (Array (SliceShape slix) e)
forall b a c d. (b -> a) -> (c -> d -> b) -> c -> d -> a
$$ FromApplyAcc
  (Acc (Array (FullShape slix) e)
   -> Exp slix -> SmartAcc (Array (EltR (SliceShape slix)) (EltR e)))
-> Acc (Array (FullShape slix) e)
-> Exp slix
-> SmartAcc (Array (EltR (SliceShape slix)) (EltR e))
forall a. ApplyAcc a => FromApplyAcc a -> a
applyAcc (SliceIndex
  (EltR slix)
  (EltR (SliceShape slix))
  (EltR (CoSliceShape slix))
  (EltR (FullShape slix))
-> SmartAcc (Array (EltR (FullShape slix)) (EltR e))
-> SmartExp (EltR slix)
-> PreSmartAcc
     SmartAcc SmartExp (Array (EltR (SliceShape slix)) (EltR e))
forall slix sl co sh (acc :: * -> *) e' (exp :: * -> *).
SliceIndex slix sl co sh
-> acc (Array sh e')
-> exp slix
-> PreSmartAcc acc exp (Array sl e')
Slice (SliceIndex
   (EltR slix)
   (EltR (SliceShape slix))
   (EltR (CoSliceShape slix))
   (EltR (FullShape slix))
 -> SmartAcc (Array (EltR (FullShape slix)) (EltR e))
 -> SmartExp (EltR slix)
 -> PreSmartAcc
      SmartAcc SmartExp (Array (EltR (SliceShape slix)) (EltR e)))
-> SliceIndex
     (EltR slix)
     (EltR (SliceShape slix))
     (EltR (CoSliceShape slix))
     (EltR (FullShape slix))
-> SmartAcc (Array (EltR (FullShape slix)) (EltR e))
-> SmartExp (EltR slix)
-> PreSmartAcc
     SmartAcc SmartExp (Array (EltR (SliceShape slix)) (EltR e))
forall a b. (a -> b) -> a -> b
$ Slice slix =>
SliceIndex
  (EltR slix)
  (EltR (SliceShape slix))
  (EltR (CoSliceShape slix))
  (EltR (FullShape slix))
forall sl.
Slice sl =>
SliceIndex
  (EltR sl)
  (EltR (SliceShape sl))
  (EltR (CoSliceShape sl))
  (EltR (FullShape sl))
sliceIndex @slix)

-- Map-like functions
-- ------------------

-- | Apply the given function element-wise to an array. Denotationally we have:
--
-- > map f [x1, x2, ... xn] = [f x1, f x2, ... f xn]
--
-- >>> let xs = fromList (Z:.10) [0..] :: Vector Int
-- >>> xs
-- Vector (Z :. 10) [0,1,2,3,4,5,6,7,8,9]
--
-- >>> run $ map (+1) (use xs)
-- Vector (Z :. 10) [1,2,3,4,5,6,7,8,9,10]
--
map :: forall sh a b.
       (Shape sh, Elt a, Elt b)
    => (Exp a -> Exp b)
    -> Acc (Array sh a)
    -> Acc (Array sh b)
map :: (Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map = SmartAcc (Array (EltR sh) (EltR b)) -> Acc (Array sh b)
forall a. SmartAcc (ArraysR a) -> Acc a
Acc (SmartAcc (Array (EltR sh) (EltR b)) -> Acc (Array sh b))
-> ((Exp a -> Exp b)
    -> Acc (Array sh a) -> SmartAcc (Array (EltR sh) (EltR b)))
-> (Exp a -> Exp b)
-> Acc (Array sh a)
-> Acc (Array sh b)
forall b a c d. (b -> a) -> (c -> d -> b) -> c -> d -> a
$$ FromApplyAcc
  ((Exp a -> Exp b)
   -> Acc (Array sh a) -> SmartAcc (Array (EltR sh) (EltR b)))
-> (Exp a -> Exp b)
-> Acc (Array sh a)
-> SmartAcc (Array (EltR sh) (EltR b))
forall a. ApplyAcc a => FromApplyAcc a -> a
applyAcc (TypeR (EltR a)
-> TypeR (EltR b)
-> (SmartExp (EltR a) -> SmartExp (EltR b))
-> SmartAcc (Array (EltR sh) (EltR a))
-> PreSmartAcc SmartAcc SmartExp (Array (EltR sh) (EltR b))
forall e e' (exp :: * -> *) (acc :: * -> *) sh.
TypeR e
-> TypeR e'
-> (SmartExp e -> exp e')
-> acc (Array sh e)
-> PreSmartAcc acc exp (Array sh e')
Map (Elt a => TypeR (EltR a)
forall a. Elt a => TypeR (EltR a)
eltR @a) (Elt b => TypeR (EltR b)
forall a. Elt a => TypeR (EltR a)
eltR @b))

-- | Apply the given binary function element-wise to the two arrays. The extent
-- of the resulting array is the intersection of the extents of the two source
-- arrays.
--
-- >>> let xs = fromList (Z:.3:.5) [0..] :: Matrix Int
-- >>> xs
-- Matrix (Z :. 3 :. 5)
--   [  0,  1,  2,  3,  4,
--      5,  6,  7,  8,  9,
--     10, 11, 12, 13, 14]
--
-- >>> let ys = fromList (Z:.5:.10) [1..] :: Matrix Int
-- >>> ys
-- Matrix (Z :. 5 :. 10)
--   [  1,  2,  3,  4,  5,  6,  7,  8,  9, 10,
--     11, 12, 13, 14, 15, 16, 17, 18, 19, 20,
--     21, 22, 23, 24, 25, 26, 27, 28, 29, 30,
--     31, 32, 33, 34, 35, 36, 37, 38, 39, 40,
--     41, 42, 43, 44, 45, 46, 47, 48, 49, 50]
--
-- >>> run $ zipWith (+) (use xs) (use ys)
-- Matrix (Z :. 3 :. 5)
--   [  1,  3,  5,  7,  9,
--     16, 18, 20, 22, 24,
--     31, 33, 35, 37, 39]
--
zipWith :: forall sh a b c.
           (Shape sh, Elt a, Elt b, Elt c)
        => (Exp a -> Exp b -> Exp c)
        -> Acc (Array sh a)
        -> Acc (Array sh b)
        -> Acc (Array sh c)
zipWith :: (Exp a -> Exp b -> Exp c)
-> Acc (Array sh a) -> Acc (Array sh b) -> Acc (Array sh c)
zipWith = SmartAcc (Array (EltR sh) (EltR c)) -> Acc (Array sh c)
forall a. SmartAcc (ArraysR a) -> Acc a
Acc (SmartAcc (Array (EltR sh) (EltR c)) -> Acc (Array sh c))
-> ((Exp a -> Exp b -> Exp c)
    -> Acc (Array sh a)
    -> Acc (Array sh b)
    -> SmartAcc (Array (EltR sh) (EltR c)))
-> (Exp a -> Exp b -> Exp c)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> Acc (Array sh c)
forall b a c d e.
(b -> a) -> (c -> d -> e -> b) -> c -> d -> e -> a
$$$ FromApplyAcc
  ((Exp a -> Exp b -> Exp c)
   -> Acc (Array sh a)
   -> Acc (Array sh b)
   -> SmartAcc (Array (EltR sh) (EltR c)))
-> (Exp a -> Exp b -> Exp c)
-> Acc (Array sh a)
-> Acc (Array sh b)
-> SmartAcc (Array (EltR sh) (EltR c))
forall a. ApplyAcc a => FromApplyAcc a -> a
applyAcc (TypeR (EltR a)
-> TypeR (EltR b)
-> TypeR (EltR c)
-> (SmartExp (EltR a) -> SmartExp (EltR b) -> SmartExp (EltR c))
-> SmartAcc (Array (EltR sh) (EltR a))
-> SmartAcc (Array (EltR sh) (EltR b))
-> PreSmartAcc SmartAcc SmartExp (Array (EltR sh) (EltR c))
forall e1 e2 e3 (exp :: * -> *) (acc :: * -> *) sh.
TypeR e1
-> TypeR e2
-> TypeR e3
-> (SmartExp e1 -> SmartExp e2 -> exp e3)
-> acc (Array sh e1)
-> acc (Array sh e2)
-> PreSmartAcc acc exp (Array sh e3)
ZipWith (Elt a => TypeR (EltR a)
forall a. Elt a => TypeR (EltR a)
eltR @a) (Elt b => TypeR (EltR b)
forall a. Elt a => TypeR (EltR a)
eltR @b) (Elt c => TypeR (EltR c)
forall a. Elt a => TypeR (EltR a)
eltR @c))

-- Reductions
-- ----------

-- | Reduction of the innermost dimension of an array of arbitrary rank.
--
-- The shape of the result obeys the property:
--
-- > shape (fold f z xs) == indexTail (shape xs)
--
-- The first argument needs to be an /associative/ function to enable an
-- efficient parallel implementation. The initial element does not need to be an
-- identity element of the combination function.
--
-- >>> let mat = fromList (Z:.5:.10) [0..] :: Matrix Int
-- >>> mat
-- Matrix (Z :. 5 :. 10)
--   [  0,  1,  2,  3,  4,  5,  6,  7,  8,  9,
--     10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
--     20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
--     30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
--     40, 41, 42, 43, 44, 45, 46, 47, 48, 49]
--
-- >>> run $ fold (+) 42 (use mat)
-- Vector (Z :. 5) [87,187,287,387,487]
--
-- Reductions with non-commutative operators are supported. For example, the
-- following computes the maximum segment sum problem along each innermost
-- dimension of the array.
--
-- <https://en.wikipedia.org/wiki/Maximum_subarray_problem>
--
-- >>> :{
--   let maximumSegmentSum
--           :: forall sh e. (Shape sh, Num e, Ord e)
--           => Acc (Array (sh :. Int) e)
--           -> Acc (Array sh e)
--       maximumSegmentSum
--         = map (\(T4 x _ _ _) -> x)
--         . fold1 f
--         . map g
--         where
--           f :: (Num a, Ord a) => Exp (a,a,a,a) -> Exp (a,a,a,a) -> Exp (a,a,a,a)
--           f x y =
--             let T4 mssx misx mcsx tsx = x
--                 T4 mssy misy mcsy tsy = y
--             in
--             T4 (mssx `max` (mssy `max` (mcsx+misy)))
--                (misx `max` (tsx+misy))
--                (mcsy `max` (mcsx+tsy))
--                (tsx+tsy)
--           --
--           g :: (Num a, Ord a) => Exp a -> Exp (a,a,a,a)
--           g x = let y = max x 0
--                  in T4 y y y x
-- :}
--
-- >>> let vec = fromList (Z:.10) [-2,1,-3,4,-1,2,1,-5,4,0] :: Vector Int
-- >>> run $ maximumSegmentSum (use vec)
-- Scalar Z [6]
--
-- See also 'Data.Array.Accelerate.Data.Fold.Fold', which can be a useful way to
-- compute multiple results from a single reduction.
--
fold :: forall sh a.
        (Shape sh, Elt a)
     => (Exp a -> Exp a -> Exp a)
     -> Exp a
     -> Acc (Array (sh:.Int) a)
     -> Acc (Array sh a)
fold :: (Exp a -> Exp a -> Exp a)
-> Exp a -> Acc (Array (sh :. Int) a) -> Acc (Array sh a)
fold Exp a -> Exp a -> Exp a
f (Exp SmartExp (EltR a)
x) = SmartAcc (Array (EltR sh) (EltR a)) -> Acc (Array sh a)
forall a. SmartAcc (ArraysR a) -> Acc a
Acc (SmartAcc (Array (EltR sh) (EltR a)) -> Acc (Array sh a))
-> (Acc (Array (sh :. Int) a)
    -> SmartAcc (Array (EltR sh) (EltR a)))
-> Acc (Array (sh :. Int) a)
-> Acc (Array sh a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FromApplyAcc
  (Acc (Array (sh :. Int) a) -> SmartAcc (Array (EltR sh) (EltR a)))
-> Acc (Array (sh :. Int) a) -> SmartAcc (Array (EltR sh) (EltR a))
forall a. ApplyAcc a => FromApplyAcc a -> a
applyAcc (TypeR (EltR a)
-> (SmartExp (EltR a) -> SmartExp (EltR a) -> SmartExp (EltR a))
-> Maybe (SmartExp (EltR a))
-> SmartAcc (Array (EltR sh, Int) (EltR a))
-> PreSmartAcc SmartAcc SmartExp (Array (EltR sh) (EltR a))
forall e (exp :: * -> *) (acc :: * -> *) i.
TypeR e
-> (SmartExp e -> SmartExp e -> exp e)
-> Maybe (exp e)
-> acc (Array (i, Int) e)
-> PreSmartAcc acc exp (Array i e)
Fold (Elt a => TypeR (EltR a)
forall a. Elt a => TypeR (EltR a)
eltR @a) ((Exp a -> Exp a -> Exp a)
-> SmartExp (EltR a) -> SmartExp (EltR a) -> SmartExp (EltR a)
forall a b c.
(Elt a, Elt b, Elt c) =>
(Exp a -> Exp b -> Exp c)
-> SmartExp (EltR a) -> SmartExp (EltR b) -> SmartExp (EltR c)
unExpBinaryFunction Exp a -> Exp a -> Exp a
f) (SmartExp (EltR a) -> Maybe (SmartExp (EltR a))
forall a. a -> Maybe a
Just SmartExp (EltR a)
x))

-- | Variant of 'fold' that requires the innermost dimension of the array to be
-- non-empty and doesn't need an default value.
--
-- The shape of the result obeys the property:
--
-- > shape (fold f z xs) == indexTail (shape xs)
--
-- The first argument needs to be an /associative/ function to enable an
-- efficient parallel implementation, but does not need to be commutative.
--
fold1 :: forall sh a.
         (Shape sh, Elt a)
      => (Exp a -> Exp a -> Exp a)
      -> Acc (Array (sh:.Int) a)
      -> Acc (Array sh a)
fold1 :: (Exp a -> Exp a -> Exp a)
-> Acc (Array (sh :. Int) a) -> Acc (Array sh a)
fold1 Exp a -> Exp a -> Exp a
f = SmartAcc (Array (EltR sh) (EltR a)) -> Acc (Array sh a)
forall a. SmartAcc (ArraysR a) -> Acc a
Acc (SmartAcc (Array (EltR sh) (EltR a)) -> Acc (Array sh a))
-> (Acc (Array (sh :. Int) a)
    -> SmartAcc (Array (EltR sh) (EltR a)))
-> Acc (Array (sh :. Int) a)
-> Acc (Array sh a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FromApplyAcc
  (Acc (Array (sh :. Int) a) -> SmartAcc (Array (EltR sh) (EltR a)))
-> Acc (Array (sh :. Int) a) -> SmartAcc (Array (EltR sh) (EltR a))
forall a. ApplyAcc a => FromApplyAcc a -> a
applyAcc (TypeR (EltR a)
-> (SmartExp (EltR a) -> SmartExp (EltR a) -> SmartExp (EltR a))
-> Maybe (SmartExp (EltR a))
-> SmartAcc (Array (EltR sh, Int) (EltR a))
-> PreSmartAcc SmartAcc SmartExp (Array (EltR sh) (EltR a))
forall e (exp :: * -> *) (acc :: * -> *) i.
TypeR e
-> (SmartExp e -> SmartExp e -> exp e)
-> Maybe (exp e)
-> acc (Array (i, Int) e)
-> PreSmartAcc acc exp (Array i e)
Fold (Elt a => TypeR (EltR a)
forall a. Elt a => TypeR (EltR a)
eltR @a) ((Exp a -> Exp a -> Exp a)
-> SmartExp (EltR a) -> SmartExp (EltR a) -> SmartExp (EltR a)
forall a b c.
(Elt a, Elt b, Elt c) =>
(Exp a -> Exp b -> Exp c)
-> SmartExp (EltR a) -> SmartExp (EltR b) -> SmartExp (EltR c)
unExpBinaryFunction Exp a -> Exp a -> Exp a
f) Maybe (SmartExp (EltR a))
forall a. Maybe a
Nothing)

-- | Segmented reduction along the innermost dimension of an array. The
-- segment descriptor specifies the starting index (offset) along the
-- innermost dimension to the beginning of each logical sub-array.
--
-- The value in the output array at index i is the reduction of values
-- between the indices of the segment descriptor at index i and (i+1).
--
-- We have that:
--
-- > foldSeg f z xs seg  ==  foldSeg' f z xs (scanl (+) 0 seg)
--
-- @since 1.3.0.0
--
foldSeg'
    :: forall sh a i.
       (Shape sh, Elt a, Elt i, IsIntegral i, i ~ EltR i)
    => (Exp a -> Exp a -> Exp a)
    -> Exp a
    -> Acc (Array (sh:.Int) a)
    -> Acc (Segments i)
    -> Acc (Array (sh:.Int) a)
foldSeg' :: (Exp a -> Exp a -> Exp a)
-> Exp a
-> Acc (Array (sh :. Int) a)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) a)
foldSeg' Exp a -> Exp a -> Exp a
f (Exp SmartExp (EltR a)
x) = SmartAcc (Array (EltR sh, Int) (EltR a))
-> Acc (Array (sh :. Int) a)
forall a. SmartAcc (ArraysR a) -> Acc a
Acc (SmartAcc (Array (EltR sh, Int) (EltR a))
 -> Acc (Array (sh :. Int) a))
-> (Acc (Array (sh :. Int) a)
    -> Acc (Segments i) -> SmartAcc (Array (EltR sh, Int) (EltR a)))
-> Acc (Array (sh :. Int) a)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) a)
forall b a c d. (b -> a) -> (c -> d -> b) -> c -> d -> a
$$ FromApplyAcc
  (Acc (Array (sh :. Int) a)
   -> Acc (Segments i) -> SmartAcc (Array (EltR sh, Int) (EltR a)))
-> Acc (Array (sh :. Int) a)
-> Acc (Segments i)
-> SmartAcc (Array (EltR sh, Int) (EltR a))
forall a. ApplyAcc a => FromApplyAcc a -> a
applyAcc (IntegralType i
-> TypeR (EltR a)
-> (SmartExp (EltR a) -> SmartExp (EltR a) -> SmartExp (EltR a))
-> Maybe (SmartExp (EltR a))
-> SmartAcc (Array (EltR sh, Int) (EltR a))
-> SmartAcc (Segments i)
-> PreSmartAcc SmartAcc SmartExp (Array (EltR sh, Int) (EltR a))
forall i e (exp :: * -> *) (acc :: * -> *) sh.
IntegralType i
-> TypeR e
-> (SmartExp e -> SmartExp e -> exp e)
-> Maybe (exp e)
-> acc (Array (sh, Int) e)
-> acc (Segments i)
-> PreSmartAcc acc exp (Array (sh, Int) e)
FoldSeg (IsIntegral i => IntegralType i
forall a. IsIntegral a => IntegralType a
integralType @i) (Elt a => TypeR (EltR a)
forall a. Elt a => TypeR (EltR a)
eltR @a) ((Exp a -> Exp a -> Exp a)
-> SmartExp (EltR a) -> SmartExp (EltR a) -> SmartExp (EltR a)
forall a b c.
(Elt a, Elt b, Elt c) =>
(Exp a -> Exp b -> Exp c)
-> SmartExp (EltR a) -> SmartExp (EltR b) -> SmartExp (EltR c)
unExpBinaryFunction Exp a -> Exp a -> Exp a
f) (SmartExp (EltR a) -> Maybe (SmartExp (EltR a))
forall a. a -> Maybe a
Just SmartExp (EltR a)
x))

-- | Variant of 'foldSeg'' that requires /all/ segments of the reduced
-- array to be non-empty, and doesn't need a default value. The segment
-- descriptor specifies the offset to the beginning of each of the logical
-- sub-arrays.
--
-- @since 1.3.0.0
--
fold1Seg'
    :: forall sh a i.
       (Shape sh, Elt a, Elt i, IsIntegral i, i ~ EltR i)
    => (Exp a -> Exp a -> Exp a)
    -> Acc (Array (sh:.Int) a)
    -> Acc (Segments i)
    -> Acc (Array (sh:.Int) a)
fold1Seg' :: (Exp a -> Exp a -> Exp a)
-> Acc (Array (sh :. Int) a)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) a)
fold1Seg' Exp a -> Exp a -> Exp a
f = SmartAcc (Array (EltR sh, Int) (EltR a))
-> Acc (Array (sh :. Int) a)
forall a. SmartAcc (ArraysR a) -> Acc a
Acc (SmartAcc (Array (EltR sh, Int) (EltR a))
 -> Acc (Array (sh :. Int) a))
-> (Acc (Array (sh :. Int) a)
    -> Acc (Segments i) -> SmartAcc (Array (EltR sh, Int) (EltR a)))
-> Acc (Array (sh :. Int) a)
-> Acc (Segments i)
-> Acc (Array (sh :. Int) a)
forall b a c d. (b -> a) -> (c -> d -> b) -> c -> d -> a
$$ FromApplyAcc
  (Acc (Array (sh :. Int) a)
   -> Acc (Segments i) -> SmartAcc (Array (EltR sh, Int) (EltR a)))
-> Acc (Array (sh :. Int) a)
-> Acc (Segments i)
-> SmartAcc (Array (EltR sh, Int) (EltR a))
forall a. ApplyAcc a => FromApplyAcc a -> a
applyAcc (IntegralType i
-> TypeR (EltR a)
-> (SmartExp (EltR a) -> SmartExp (EltR a) -> SmartExp (EltR a))
-> Maybe (SmartExp (EltR a))
-> SmartAcc (Array (EltR sh, Int) (EltR a))
-> SmartAcc (Segments i)
-> PreSmartAcc SmartAcc SmartExp (Array (EltR sh, Int) (EltR a))
forall i e (exp :: * -> *) (acc :: * -> *) sh.
IntegralType i
-> TypeR e
-> (SmartExp e -> SmartExp e -> exp e)
-> Maybe (exp e)
-> acc (Array (sh, Int) e)
-> acc (Segments i)
-> PreSmartAcc acc exp (Array (sh, Int) e)
FoldSeg (IsIntegral i => IntegralType i
forall a. IsIntegral a => IntegralType a
integralType @i) (Elt a => TypeR (EltR a)
forall a. Elt a => TypeR (EltR a)
eltR @a) ((Exp a -> Exp a -> Exp a)
-> SmartExp (EltR a) -> SmartExp (EltR a) -> SmartExp (EltR a)
forall a b c.
(Elt a, Elt b, Elt c) =>
(Exp a -> Exp b -> Exp c)
-> SmartExp (EltR a) -> SmartExp (EltR b) -> SmartExp (EltR c)
unExpBinaryFunction Exp a -> Exp a -> Exp a
f) Maybe (SmartExp (EltR a))
forall a. Maybe a
Nothing)

-- Scan functions
-- --------------

-- | Data.List style left-to-right scan along the innermost dimension of an
-- arbitrary rank array. The first argument needs to be an /associative/
-- function to enable efficient parallel implementation. The initial value
-- (second argument) may be arbitrary.
--
-- >>> let vec = fromList (Z :. 10) [0..] :: Vector Int
-- >>> run $ scanl (+) 10 (use vec)
-- Vector (Z :. 11) [10,10,11,13,16,20,25,31,38,46,55]
--
-- >>> let mat = fromList (Z :. 4 :. 10) [0..] :: Matrix Int
-- >>> run $ scanl (+) 0 (use mat)
-- Matrix (Z :. 4 :. 11)
--   [ 0,  0,  1,  3,   6,  10,  15,  21,  28,  36,  45,
--     0, 10, 21, 33,  46,  60,  75,  91, 108, 126, 145,
--     0, 20, 41, 63,  86, 110, 135, 161, 188, 216, 245,
--     0, 30, 61, 93, 126, 160, 195, 231, 268, 306, 345]
--
scanl :: forall sh a.
         (Shape sh, Elt a)
      => (Exp a -> Exp a -> Exp a)
      -> Exp a
      -> Acc (Array (sh:.Int) a)
      -> Acc (Array (sh:.Int) a)
scanl :: (Exp a -> Exp a -> Exp a)
-> Exp a -> Acc (Array (sh :. Int) a) -> Acc (Array (sh :. Int) a)
scanl Exp a -> Exp a -> Exp a
f (Exp SmartExp (EltR a)
x) (Acc SmartAcc (ArraysR (Array (sh :. Int) a))
a) = SmartAcc (ArraysR (Array (sh :. Int) a))
-> Acc (Array (sh :. Int) a)
forall a. SmartAcc (ArraysR a) -> Acc a
Acc (SmartAcc (ArraysR (Array (sh :. Int) a))
 -> Acc (Array (sh :. Int) a))
-> SmartAcc (ArraysR (Array (sh :. Int) a))
-> Acc (Array (sh :. Int) a)
forall a b. (a -> b) -> a -> b
$ PreSmartAcc SmartAcc SmartExp (Array (EltR sh, Int) (EltR a))
-> SmartAcc (Array (EltR sh, Int) (EltR a))
forall a. PreSmartAcc SmartAcc SmartExp a -> SmartAcc a
SmartAcc (PreSmartAcc SmartAcc SmartExp (Array (EltR sh, Int) (EltR a))
 -> SmartAcc (Array (EltR sh, Int) (EltR a)))
-> PreSmartAcc SmartAcc SmartExp (Array (EltR sh, Int) (EltR a))
-> SmartAcc (Array (EltR sh, Int) (EltR a))
forall a b. (a -> b) -> a -> b
$ Direction
-> TypeR (EltR a)
-> (SmartExp (EltR a) -> SmartExp (EltR a) -> SmartExp (EltR a))
-> Maybe (SmartExp (EltR a))
-> SmartAcc (Array (EltR sh, Int) (EltR a))
-> PreSmartAcc SmartAcc SmartExp (Array (EltR sh, Int) (EltR a))
forall e (exp :: * -> *) (acc :: * -> *) e.
Direction
-> TypeR e
-> (SmartExp e -> SmartExp e -> exp e)
-> Maybe (exp e)
-> acc (Array (e, Int) e)
-> PreSmartAcc acc exp (Array (e, Int) e)
Scan Direction
LeftToRight (Elt a => TypeR (EltR a)
forall a. Elt a => TypeR (EltR a)
eltR @a) ((Exp a -> Exp a -> Exp a)
-> SmartExp (EltR a) -> SmartExp (EltR a) -> SmartExp (EltR a)
forall a b c.
(Elt a, Elt b, Elt c) =>
(Exp a -> Exp b -> Exp c)
-> SmartExp (EltR a) -> SmartExp (EltR b) -> SmartExp (EltR c)
unExpBinaryFunction Exp a -> Exp a -> Exp a
f) (SmartExp (EltR a) -> Maybe (SmartExp (EltR a))
forall a. a -> Maybe a
Just SmartExp (EltR a)
x) SmartAcc (Array (EltR sh, Int) (EltR a))
SmartAcc (ArraysR (Array (sh :. Int) a))
a

-- | Variant of 'scanl', where the last element (final reduction result) along
-- each dimension is returned separately. Denotationally we have:
--
-- > scanl' f e arr = (init res, unit (res!len))
-- >   where
-- >     len = shape arr
-- >     res = scanl f e arr
--
-- >>> let vec       = fromList (Z:.10) [0..] :: Vector Int
-- >>> let (res,sum) = run $ scanl' (+) 0 (use vec)
-- >>> res
-- Vector (Z :. 10) [0,0,1,3,6,10,15,21,28,36]
-- >>> sum
-- Scalar Z [45]
--
-- >>> let mat        = fromList (Z:.4:.10) [0..] :: Matrix Int
-- >>> let (res,sums) = run $ scanl' (+) 0 (use mat)
-- >>> res
-- Matrix (Z :. 4 :. 10)
--   [ 0,  0,  1,  3,   6,  10,  15,  21,  28,  36,
--     0, 10, 21, 33,  46,  60,  75,  91, 108, 126,
--     0, 20, 41, 63,  86, 110, 135, 161, 188, 216,
--     0, 30, 61, 93, 126, 160, 195, 231, 268, 306]
-- >>> sums
-- Vector (Z :. 4) [45,145,245,345]
--
scanl' :: forall sh a.
          (Shape sh, Elt a)
       => (Exp a -> Exp a -> Exp a)
       -> Exp a
       -> Acc (Array (sh:.Int) a)
       -> Acc (Array (sh:.Int) a, Array sh a)
scanl' :: (Exp a -> Exp a -> Exp a)
-> Exp a
-> Acc (Array (sh :. Int) a)
-> Acc (Array (sh :. Int) a, Array sh a)
scanl' = SmartAcc
  (((), Array (EltR sh, Int) (EltR a)), Array (EltR sh) (EltR a))
-> Acc (Array (sh :. Int) a, Array sh a)
forall a. SmartAcc (ArraysR a) -> Acc a
Acc (SmartAcc
   (((), Array (EltR sh, Int) (EltR a)), Array (EltR sh) (EltR a))
 -> Acc (Array (sh :. Int) a, Array sh a))
-> (SmartAcc
      (Array (EltR sh, Int) (EltR a), Array (EltR sh) (EltR a))
    -> SmartAcc
         (((), Array (EltR sh, Int) (EltR a)), Array (EltR sh) (EltR a)))
-> SmartAcc
     (Array (EltR sh, Int) (EltR a), Array (EltR sh) (EltR a))
-> Acc (Array (sh :. Int) a, Array sh a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SmartAcc (Array (EltR sh, Int) (EltR a), Array (EltR sh) (EltR a))
-> SmartAcc
     (((), Array (EltR sh, Int) (EltR a)), Array (EltR sh) (EltR a))
forall a b. SmartAcc (a, b) -> SmartAcc (((), a), b)
mkPairToTuple (SmartAcc (Array (EltR sh, Int) (EltR a), Array (EltR sh) (EltR a))
 -> Acc (Array (sh :. Int) a, Array sh a))
-> ((Exp a -> Exp a -> Exp a)
    -> Exp a
    -> Acc (Array (sh :. Int) a)
    -> SmartAcc
         (Array (EltR sh, Int) (EltR a), Array (EltR sh) (EltR a)))
-> (Exp a -> Exp a -> Exp a)
-> Exp a
-> Acc (Array (sh :. Int) a)
-> Acc (Array (sh :. Int) a, Array sh a)
forall b a c d e.
(b -> a) -> (c -> d -> e -> b) -> c -> d -> e -> a
$$$ FromApplyAcc
  ((Exp a -> Exp a -> Exp a)
   -> Exp a
   -> Acc (Array (sh :. Int) a)
   -> SmartAcc
        (Array (EltR sh, Int) (EltR a), Array (EltR sh) (EltR a)))
-> (Exp a -> Exp a -> Exp a)
-> Exp a
-> Acc (Array (sh :. Int) a)
-> SmartAcc
     (Array (EltR sh, Int) (EltR a), Array (EltR sh) (EltR a))
forall a. ApplyAcc a => FromApplyAcc a -> a
applyAcc (Direction
-> TypeR (EltR a)
-> (SmartExp (EltR a) -> SmartExp (EltR a) -> SmartExp (EltR a))
-> SmartExp (EltR a)
-> SmartAcc (Array (EltR sh, Int) (EltR a))
-> PreSmartAcc
     SmartAcc
     SmartExp
     (Array (EltR sh, Int) (EltR a), Array (EltR sh) (EltR a))
forall e (exp :: * -> *) (acc :: * -> *) sh.
Direction
-> TypeR e
-> (SmartExp e -> SmartExp e -> exp e)
-> exp e
-> acc (Array (sh, Int) e)
-> PreSmartAcc acc exp (Array (sh, Int) e, Array sh e)
Scan' Direction
LeftToRight (TypeR (EltR a)
 -> (SmartExp (EltR a) -> SmartExp (EltR a) -> SmartExp (EltR a))
 -> SmartExp (EltR a)
 -> SmartAcc (Array (EltR sh, Int) (EltR a))
 -> PreSmartAcc
      SmartAcc
      SmartExp
      (Array (EltR sh, Int) (EltR a), Array (EltR sh) (EltR a)))
-> TypeR (EltR a)
-> (SmartExp (EltR a) -> SmartExp (EltR a) -> SmartExp (EltR a))
-> SmartExp (EltR a)
-> SmartAcc (Array (EltR sh, Int) (EltR a))
-> PreSmartAcc
     SmartAcc
     SmartExp
     (Array (EltR sh, Int) (EltR a), Array (EltR sh) (EltR a))
forall a b. (a -> b) -> a -> b
$ Elt a => TypeR (EltR a)
forall a. Elt a => TypeR (EltR a)
eltR @a)

-- | Data.List style left-to-right scan along the innermost dimension without an
-- initial value (aka inclusive scan). The innermost dimension of the array must
-- not be empty. The first argument must be an /associative/ function.
--
-- >>> let mat = fromList (Z:.4:.10) [0..] :: Matrix Int
-- >>> run $ scanl1 (+) (use mat)
-- Matrix (Z :. 4 :. 10)
--   [  0,  1,  3,   6,  10,  15,  21,  28,  36,  45,
--     10, 21, 33,  46,  60,  75,  91, 108, 126, 145,
--     20, 41, 63,  86, 110, 135, 161, 188, 216, 245,
--     30, 61, 93, 126, 160, 195, 231, 268, 306, 345]
--
scanl1 :: forall sh a.
          (Shape sh, Elt a)
       => (Exp a -> Exp a -> Exp a)
       -> Acc (Array (sh:.Int) a)
       -> Acc (Array (sh:.Int) a)
scanl1 :: (Exp a -> Exp a -> Exp a)
-> Acc (Array (sh :. Int) a) -> Acc (Array (sh :. Int) a)
scanl1 Exp a -> Exp a -> Exp a
f (Acc SmartAcc (ArraysR (Array (sh :. Int) a))
a) = SmartAcc (ArraysR (Array (sh :. Int) a))
-> Acc (Array (sh :. Int) a)
forall a. SmartAcc (ArraysR a) -> Acc a
Acc (SmartAcc (ArraysR (Array (sh :. Int) a))
 -> Acc (Array (sh :. Int) a))
-> SmartAcc (ArraysR (Array (sh :. Int) a))
-> Acc (Array (sh :. Int) a)
forall a b. (a -> b) -> a -> b
$ PreSmartAcc SmartAcc SmartExp (Array (EltR sh, Int) (EltR a))
-> SmartAcc (Array (EltR sh, Int) (EltR a))
forall a. PreSmartAcc SmartAcc SmartExp a -> SmartAcc a
SmartAcc (PreSmartAcc SmartAcc SmartExp (Array (EltR sh, Int) (EltR a))
 -> SmartAcc (Array (EltR sh, Int) (EltR a)))
-> PreSmartAcc SmartAcc SmartExp (Array (EltR sh, Int) (EltR a))
-> SmartAcc (Array (EltR sh, Int) (EltR a))
forall a b. (a -> b) -> a -> b
$ Direction
-> TypeR (EltR a)
-> (SmartExp (EltR a) -> SmartExp (EltR a) -> SmartExp (EltR a))
-> Maybe (SmartExp (EltR a))
-> SmartAcc (Array (EltR sh, Int) (EltR a))
-> PreSmartAcc SmartAcc SmartExp (Array (EltR sh, Int) (EltR a))
forall e (exp :: * -> *) (acc :: * -> *) e.
Direction
-> TypeR e
-> (SmartExp e -> SmartExp e -> exp e)
-> Maybe (exp e)
-> acc (Array (e, Int) e)
-> PreSmartAcc acc exp (Array (e, Int) e)
Scan Direction
LeftToRight (Elt a => TypeR (EltR a)
forall a. Elt a => TypeR (EltR a)
eltR @a) ((Exp a -> Exp a -> Exp a)
-> SmartExp (EltR a) -> SmartExp (EltR a) -> SmartExp (EltR a)
forall a b c.
(Elt a, Elt b, Elt c) =>
(Exp a -> Exp b -> Exp c)
-> SmartExp (EltR a) -> SmartExp (EltR b) -> SmartExp (EltR c)
unExpBinaryFunction Exp a -> Exp a -> Exp a
f) Maybe (SmartExp (EltR a))
forall a. Maybe a
Nothing SmartAcc (Array (EltR sh, Int) (EltR a))
SmartAcc (ArraysR (Array (sh :. Int) a))
a

-- | Right-to-left variant of 'scanl'.
--
scanr :: forall sh a.
         (Shape sh, Elt a)
      => (Exp a -> Exp a -> Exp a)
      -> Exp a
      -> Acc (Array (sh:.Int) a)
      -> Acc (Array (sh:.Int) a)
scanr :: (Exp a -> Exp a -> Exp a)
-> Exp a -> Acc (Array (sh :. Int) a) -> Acc (Array (sh :. Int) a)
scanr Exp a -> Exp a -> Exp a
f (Exp SmartExp (EltR a)
x) (Acc SmartAcc (ArraysR (Array (sh :. Int) a))
a) = SmartAcc (ArraysR (Array (sh :. Int) a))
-> Acc (Array (sh :. Int) a)
forall a. SmartAcc (ArraysR a) -> Acc a
Acc (SmartAcc (ArraysR (Array (sh :. Int) a))
 -> Acc (Array (sh :. Int) a))
-> SmartAcc (ArraysR (Array (sh :. Int) a))
-> Acc (Array (sh :. Int) a)
forall a b. (a -> b) -> a -> b
$ PreSmartAcc SmartAcc SmartExp (Array (EltR sh, Int) (EltR a))
-> SmartAcc (Array (EltR sh, Int) (EltR a))
forall a. PreSmartAcc SmartAcc SmartExp a -> SmartAcc a
SmartAcc (PreSmartAcc SmartAcc SmartExp (Array (EltR sh, Int) (EltR a))
 -> SmartAcc (Array (EltR sh, Int) (EltR a)))
-> PreSmartAcc SmartAcc SmartExp (Array (EltR sh, Int) (EltR a))
-> SmartAcc (Array (EltR sh, Int) (EltR a))
forall a b. (a -> b) -> a -> b
$ Direction
-> TypeR (EltR a)
-> (SmartExp (EltR a) -> SmartExp (EltR a) -> SmartExp (EltR a))
-> Maybe (SmartExp (EltR a))
-> SmartAcc (Array (EltR sh, Int) (EltR a))
-> PreSmartAcc SmartAcc SmartExp (Array (EltR sh, Int) (EltR a))
forall e (exp :: * -> *) (acc :: * -> *) e.
Direction
-> TypeR e
-> (SmartExp e -> SmartExp e -> exp e)
-> Maybe (exp e)
-> acc (Array (e, Int) e)
-> PreSmartAcc acc exp (Array (e, Int) e)
Scan Direction
RightToLeft (Elt a => TypeR (EltR a)
forall a. Elt a => TypeR (EltR a)
eltR @a) ((Exp a -> Exp a -> Exp a)
-> SmartExp (EltR a) -> SmartExp (EltR a) -> SmartExp (EltR a)
forall a b c.
(Elt a, Elt b, Elt c) =>
(Exp a -> Exp b -> Exp c)
-> SmartExp (EltR a) -> SmartExp (EltR b) -> SmartExp (EltR c)
unExpBinaryFunction Exp a -> Exp a -> Exp a
f) (SmartExp (EltR a) -> Maybe (SmartExp (EltR a))
forall a. a -> Maybe a
Just SmartExp (EltR a)
x) SmartAcc (Array (EltR sh, Int) (EltR a))
SmartAcc (ArraysR (Array (sh :. Int) a))
a

-- | Right-to-left variant of 'scanl''.
--
scanr' :: forall sh a.
          (Shape sh, Elt a)
       => (Exp a -> Exp a -> Exp a)
       -> Exp a
       -> Acc (Array (sh:.Int) a)
       -> Acc (Array (sh:.Int) a, Array sh a)
scanr' :: (Exp a -> Exp a -> Exp a)
-> Exp a
-> Acc (Array (sh :. Int) a)
-> Acc (Array (sh :. Int) a, Array sh a)
scanr' = SmartAcc
  (((), Array (EltR sh, Int) (EltR a)), Array (EltR sh) (EltR a))
-> Acc (Array (sh :. Int) a, Array sh a)
forall a. SmartAcc (ArraysR a) -> Acc a
Acc (SmartAcc
   (((), Array (EltR sh, Int) (EltR a)), Array (EltR sh) (EltR a))
 -> Acc (Array (sh :. Int) a, Array sh a))
-> (SmartAcc
      (Array (EltR sh, Int) (EltR a), Array (EltR sh) (EltR a))
    -> SmartAcc
         (((), Array (EltR sh, Int) (EltR a)), Array (EltR sh) (EltR a)))
-> SmartAcc
     (Array (EltR sh, Int) (EltR a), Array (EltR sh) (EltR a))
-> Acc (Array (sh :. Int) a, Array sh a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SmartAcc (Array (EltR sh, Int) (EltR a), Array (EltR sh) (EltR a))
-> SmartAcc
     (((), Array (EltR sh, Int) (EltR a)), Array (EltR sh) (EltR a))
forall a b. SmartAcc (a, b) -> SmartAcc (((), a), b)
mkPairToTuple (SmartAcc (Array (EltR sh, Int) (EltR a), Array (EltR sh) (EltR a))
 -> Acc (Array (sh :. Int) a, Array sh a))
-> ((Exp a -> Exp a -> Exp a)
    -> Exp a
    -> Acc (Array (sh :. Int) a)
    -> SmartAcc
         (Array (EltR sh, Int) (EltR a), Array (EltR sh) (EltR a)))
-> (Exp a -> Exp a -> Exp a)
-> Exp a
-> Acc (Array (sh :. Int) a)
-> Acc (Array (sh :. Int) a, Array sh a)
forall b a c d e.
(b -> a) -> (c -> d -> e -> b) -> c -> d -> e -> a
$$$ FromApplyAcc
  ((Exp a -> Exp a -> Exp a)
   -> Exp a
   -> Acc (Array (sh :. Int) a)
   -> SmartAcc
        (Array (EltR sh, Int) (EltR a), Array (EltR sh) (EltR a)))
-> (Exp a -> Exp a -> Exp a)
-> Exp a
-> Acc (Array (sh :. Int) a)
-> SmartAcc
     (Array (EltR sh, Int) (EltR a), Array (EltR sh) (EltR a))
forall a. ApplyAcc a => FromApplyAcc a -> a
applyAcc (Direction
-> TypeR (EltR a)
-> (SmartExp (EltR a) -> SmartExp (EltR a) -> SmartExp (EltR a))
-> SmartExp (EltR a)
-> SmartAcc (Array (EltR sh, Int) (EltR a))
-> PreSmartAcc
     SmartAcc
     SmartExp
     (Array (EltR sh, Int) (EltR a), Array (EltR sh) (EltR a))
forall e (exp :: * -> *) (acc :: * -> *) sh.
Direction
-> TypeR e
-> (SmartExp e -> SmartExp e -> exp e)
-> exp e
-> acc (Array (sh, Int) e)
-> PreSmartAcc acc exp (Array (sh, Int) e, Array sh e)
Scan' Direction
RightToLeft (TypeR (EltR a)
 -> (SmartExp (EltR a) -> SmartExp (EltR a) -> SmartExp (EltR a))
 -> SmartExp (EltR a)
 -> SmartAcc (Array (EltR sh, Int) (EltR a))
 -> PreSmartAcc
      SmartAcc
      SmartExp
      (Array (EltR sh, Int) (EltR a), Array (EltR sh) (EltR a)))
-> TypeR (EltR a)
-> (SmartExp (EltR a) -> SmartExp (EltR a) -> SmartExp (EltR a))
-> SmartExp (EltR a)
-> SmartAcc (Array (EltR sh, Int) (EltR a))
-> PreSmartAcc
     SmartAcc
     SmartExp
     (Array (EltR sh, Int) (EltR a), Array (EltR sh) (EltR a))
forall a b. (a -> b) -> a -> b
$ Elt a => TypeR (EltR a)
forall a. Elt a => TypeR (EltR a)
eltR @a)

-- | Right-to-left variant of 'scanl1'.
--
scanr1 :: forall sh a.
          (Shape sh, Elt a)
       => (Exp a -> Exp a -> Exp a)
       -> Acc (Array (sh:.Int) a)
       -> Acc (Array (sh:.Int) a)
scanr1 :: (Exp a -> Exp a -> Exp a)
-> Acc (Array (sh :. Int) a) -> Acc (Array (sh :. Int) a)
scanr1 Exp a -> Exp a -> Exp a
f (Acc SmartAcc (ArraysR (Array (sh :. Int) a))
a) = SmartAcc (ArraysR (Array (sh :. Int) a))
-> Acc (Array (sh :. Int) a)
forall a. SmartAcc (ArraysR a) -> Acc a
Acc (SmartAcc (ArraysR (Array (sh :. Int) a))
 -> Acc (Array (sh :. Int) a))
-> SmartAcc (ArraysR (Array (sh :. Int) a))
-> Acc (Array (sh :. Int) a)
forall a b. (a -> b) -> a -> b
$ PreSmartAcc SmartAcc SmartExp (Array (EltR sh, Int) (EltR a))
-> SmartAcc (Array (EltR sh, Int) (EltR a))
forall a. PreSmartAcc SmartAcc SmartExp a -> SmartAcc a
SmartAcc (PreSmartAcc SmartAcc SmartExp (Array (EltR sh, Int) (EltR a))
 -> SmartAcc (Array (EltR sh, Int) (EltR a)))
-> PreSmartAcc SmartAcc SmartExp (Array (EltR sh, Int) (EltR a))
-> SmartAcc (Array (EltR sh, Int) (EltR a))
forall a b. (a -> b) -> a -> b
$ Direction
-> TypeR (EltR a)
-> (SmartExp (EltR a) -> SmartExp (EltR a) -> SmartExp (EltR a))
-> Maybe (SmartExp (EltR a))
-> SmartAcc (Array (EltR sh, Int) (EltR a))
-> PreSmartAcc SmartAcc SmartExp (Array (EltR sh, Int) (EltR a))
forall e (exp :: * -> *) (acc :: * -> *) e.
Direction
-> TypeR e
-> (SmartExp e -> SmartExp e -> exp e)
-> Maybe (exp e)
-> acc (Array (e, Int) e)
-> PreSmartAcc acc exp (Array (e, Int) e)
Scan Direction
RightToLeft (Elt a => TypeR (EltR a)
forall a. Elt a => TypeR (EltR a)
eltR @a) ((Exp a -> Exp a -> Exp a)
-> SmartExp (EltR a) -> SmartExp (EltR a) -> SmartExp (EltR a)
forall a b c.
(Elt a, Elt b, Elt c) =>
(Exp a -> Exp b -> Exp c)
-> SmartExp (EltR a) -> SmartExp (EltR b) -> SmartExp (EltR c)
unExpBinaryFunction Exp a -> Exp a -> Exp a
f) Maybe (SmartExp (EltR a))
forall a. Maybe a
Nothing SmartAcc (Array (EltR sh, Int) (EltR a))
SmartAcc (ArraysR (Array (sh :. Int) a))
a

-- Permutations
-- ------------

-- | Generalised forward permutation operation (array scatter).
--
-- Forward permutation specified by a function mapping indices from the source
-- array to indices in the result array. The result array is initialised with
-- the given defaults and any further values that are permuted into the result
-- array are added to the current value using the given combination function.
--
-- The combination function must be /associative/ and /commutative/.
-- Elements for which the permutation function returns 'Nothing' are
-- dropped.
--
-- The combination function is given the new value being permuted as its first
-- argument, and the current value of the array as its second.
--
-- For example, we can use 'permute' to compute the occurrence count (histogram)
-- for an array of values in the range @[0,10)@:
--
-- >>> :{
--   let histogram :: Acc (Vector Int) -> Acc (Vector Int)
--       histogram xs =
--         let zeros = fill (constant (Z:.10)) 0
--             ones  = fill (shape xs)         1
--         in
--         permute (+) zeros (\ix -> Just_ (I1 (xs!ix))) ones
-- :}
--
-- >>> let xs = fromList (Z :. 20) [0,0,1,2,1,1,2,4,8,3,4,9,8,3,2,5,5,3,1,2] :: Vector Int
-- >>> run $ histogram (use xs)
-- Vector (Z :. 10) [2,4,4,3,2,2,0,0,2,1]
--
-- As a second example, note that the dimensionality of the source and
-- destination arrays can differ. In this way, we can use 'permute' to create an
-- identity matrix by overwriting elements along the diagonal:
--
-- >>> :{
--   let identity :: Num a => Exp Int -> Acc (Matrix a)
--       identity n =
--         let zeros = fill (I2 n n) 0
--             ones  = fill (I1 n)   1
--         in
--         permute const zeros (\(I1 i) -> Just_ (I2 i i)) ones
-- :}
--
-- >>> run $ identity 5 :: Matrix Int
-- Matrix (Z :. 5 :. 5)
--   [ 1, 0, 0, 0, 0,
--     0, 1, 0, 0, 0,
--     0, 0, 1, 0, 0,
--     0, 0, 0, 1, 0,
--     0, 0, 0, 0, 1]
--
-- [/Note:/]
--
-- Regarding array fusion:
--
--   1. The 'permute' operation will always be evaluated; it can not be fused
--      into a later step.
--
--   2. Since the index permutation function might not cover all positions in
--      the output array (the function is not surjective), the array of default
--      values must be evaluated. However, other operations may fuse into this.
--
--   3. The array of source values can fuse into the permutation operation.
--
--   4. If the array of default values is only used once, it will be updated
--      in-place. This behaviour can be disabled this with @-fno-inplace@.
--
-- Regarding the defaults array:
--
-- If you are sure that the default values are not necessary---they are not used
-- by the combination function and every element will be overwritten---a default
-- array created by 'Data.Array.Accelerate.Prelude.fill'ing with the value
-- 'Data.Array.Accelerate.Unsafe.undef' will give you a new uninitialised array.
--
-- Regarding the combination function:
--
-- The function 'const' can be used to replace elements of the defaults
-- array with the new values. If the permutation function maps multiple
-- values to the same location in the results array (the function is not
-- injective) then this operation is non-deterministic.
--
-- Since Accelerate uses an unzipped struct-of-array representation, where
-- the individual components of product types (for example, pairs) are
-- stored in separate arrays, storing values of product type requires
-- multiple store instructions.
--
-- Accelerate prior to version 1.3.0.0 performs this operation atomically,
-- to ensure that the stored values are always consistent (each component
-- of the product type is written by the same thread). Later versions relax
-- this restriction, but this behaviour can be disabled with
-- @-fno-fast-permute-const@.
--
permute
    :: forall sh sh' a. (Shape sh, Shape sh', Elt a)
    => (Exp a -> Exp a -> Exp a)        -- ^ combination function
    -> Acc (Array sh' a)                -- ^ array of default values
    -> (Exp sh -> Exp (Maybe sh'))      -- ^ index permutation function
    -> Acc (Array sh  a)                -- ^ array of source values to be permuted
    -> Acc (Array sh' a)
permute :: (Exp a -> Exp a -> Exp a)
-> Acc (Array sh' a)
-> (Exp sh -> Exp (Maybe sh'))
-> Acc (Array sh a)
-> Acc (Array sh' a)
permute = SmartAcc (Array (EltR sh') (EltR a)) -> Acc (Array sh' a)
forall a. SmartAcc (ArraysR a) -> Acc a
Acc (SmartAcc (Array (EltR sh') (EltR a)) -> Acc (Array sh' a))
-> ((Exp a -> Exp a -> Exp a)
    -> Acc (Array sh' a)
    -> (Exp sh -> Exp (Maybe sh'))
    -> Acc (Array sh a)
    -> SmartAcc (Array (EltR sh') (EltR a)))
-> (Exp a -> Exp a -> Exp a)
-> Acc (Array sh' a)
-> (Exp sh -> Exp (Maybe sh'))
-> Acc (Array sh a)
-> Acc (Array sh' a)
forall b a c d e f.
(b -> a) -> (c -> d -> e -> f -> b) -> c -> d -> e -> f -> a
$$$$ FromApplyAcc
  ((Exp a -> Exp a -> Exp a)
   -> Acc (Array sh' a)
   -> (Exp sh -> Exp (Maybe sh'))
   -> Acc (Array sh a)
   -> SmartAcc (Array (EltR sh') (EltR a)))
-> (Exp a -> Exp a -> Exp a)
-> Acc (Array sh' a)
-> (Exp sh -> Exp (Maybe sh'))
-> Acc (Array sh a)
-> SmartAcc (Array (EltR sh') (EltR a))
forall a. ApplyAcc a => FromApplyAcc a -> a
applyAcc (ArrayR (Array (EltR sh) (EltR a))
-> (SmartExp (EltR a) -> SmartExp (EltR a) -> SmartExp (EltR a))
-> SmartAcc (Array (EltR sh') (EltR a))
-> (SmartExp (EltR sh) -> SmartExp (PrimMaybe (EltR sh')))
-> SmartAcc (Array (EltR sh) (EltR a))
-> PreSmartAcc SmartAcc SmartExp (Array (EltR sh') (EltR a))
forall sh e (exp :: * -> *) (acc :: * -> *) sh'.
ArrayR (Array sh e)
-> (SmartExp e -> SmartExp e -> exp e)
-> acc (Array sh' e)
-> (SmartExp sh -> exp (PrimMaybe sh'))
-> acc (Array sh e)
-> PreSmartAcc acc exp (Array sh' e)
Permute (ArrayR (Array (EltR sh) (EltR a))
 -> (SmartExp (EltR a) -> SmartExp (EltR a) -> SmartExp (EltR a))
 -> SmartAcc (Array (EltR sh') (EltR a))
 -> (SmartExp (EltR sh) -> SmartExp (PrimMaybe (EltR sh')))
 -> SmartAcc (Array (EltR sh) (EltR a))
 -> PreSmartAcc SmartAcc SmartExp (Array (EltR sh') (EltR a)))
-> ArrayR (Array (EltR sh) (EltR a))
-> (SmartExp (EltR a) -> SmartExp (EltR a) -> SmartExp (EltR a))
-> SmartAcc (Array (EltR sh') (EltR a))
-> (SmartExp (EltR sh) -> SmartExp (PrimMaybe (EltR sh')))
-> SmartAcc (Array (EltR sh) (EltR a))
-> PreSmartAcc SmartAcc SmartExp (Array (EltR sh') (EltR a))
forall a b. (a -> b) -> a -> b
$ (Shape sh, Elt a) => ArrayR (Array (EltR sh) (EltR a))
forall sh e. (Shape sh, Elt e) => ArrayR (Array (EltR sh) (EltR e))
arrayR @sh @a)

-- | Generalised backward permutation operation (array gather).
--
-- Backward permutation specified by a function mapping indices in the
-- destination array to indices in the source array. Elements of the output
-- array are thus generated by reading from the corresponding index in the
-- source array.
--
-- For example, backpermute can be used to
-- 'Data.Array.Accelerate.Prelude.transpose' a matrix; at every index @Z:.y:.x@
-- in the result array, we get the value at that index by reading from the
-- source array at index @Z:.x:.y@:
--
-- >>> :{
--   let swap :: Exp DIM2 -> Exp DIM2
--       swap = lift1 f
--         where
--           f :: Z :. Exp Int :. Exp Int -> Z :. Exp Int :. Exp Int
--           f (Z:.y:.x) = Z :. x :. y
-- :}
--
-- >>> let mat = fromList (Z:.5:.10) [0..] :: Matrix Int
-- >>> mat
-- Matrix (Z :. 5 :. 10)
--   [  0,  1,  2,  3,  4,  5,  6,  7,  8,  9,
--     10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
--     20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
--     30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
--     40, 41, 42, 43, 44, 45, 46, 47, 48, 49]
--
-- >>> let mat' = use mat
-- >>> run $ backpermute (swap (shape mat')) swap mat'
-- Matrix (Z :. 10 :. 5)
--   [ 0, 10, 20, 30, 40,
--     1, 11, 21, 31, 41,
--     2, 12, 22, 32, 42,
--     3, 13, 23, 33, 43,
--     4, 14, 24, 34, 44,
--     5, 15, 25, 35, 45,
--     6, 16, 26, 36, 46,
--     7, 17, 27, 37, 47,
--     8, 18, 28, 38, 48,
--     9, 19, 29, 39, 49]
--
backpermute
    :: forall sh sh' a. (Shape sh, Shape sh', Elt a)
    => Exp sh'                          -- ^ shape of the result array
    -> (Exp sh' -> Exp sh)              -- ^ index permutation function
    -> Acc (Array sh  a)                -- ^ source array
    -> Acc (Array sh' a)
backpermute :: Exp sh'
-> (Exp sh' -> Exp sh) -> Acc (Array sh a) -> Acc (Array sh' a)
backpermute = SmartAcc (Array (EltR sh') (EltR a)) -> Acc (Array sh' a)
forall a. SmartAcc (ArraysR a) -> Acc a
Acc (SmartAcc (Array (EltR sh') (EltR a)) -> Acc (Array sh' a))
-> (Exp sh'
    -> (Exp sh' -> Exp sh)
    -> Acc (Array sh a)
    -> SmartAcc (Array (EltR sh') (EltR a)))
-> Exp sh'
-> (Exp sh' -> Exp sh)
-> Acc (Array sh a)
-> Acc (Array sh' a)
forall b a c d e.
(b -> a) -> (c -> d -> e -> b) -> c -> d -> e -> a
$$$ FromApplyAcc
  (Exp sh'
   -> (Exp sh' -> Exp sh)
   -> Acc (Array sh a)
   -> SmartAcc (Array (EltR sh') (EltR a)))
-> Exp sh'
-> (Exp sh' -> Exp sh)
-> Acc (Array sh a)
-> SmartAcc (Array (EltR sh') (EltR a))
forall a. ApplyAcc a => FromApplyAcc a -> a
applyAcc (ShapeR (EltR sh')
-> SmartExp (EltR sh')
-> (SmartExp (EltR sh') -> SmartExp (EltR sh))
-> SmartAcc (Array (EltR sh) (EltR a))
-> PreSmartAcc SmartAcc SmartExp (Array (EltR sh') (EltR a))
forall sh' (exp :: * -> *) sh (acc :: * -> *) e.
ShapeR sh'
-> exp sh'
-> (SmartExp sh' -> exp sh)
-> acc (Array sh e)
-> PreSmartAcc acc exp (Array sh' e)
Backpermute (ShapeR (EltR sh')
 -> SmartExp (EltR sh')
 -> (SmartExp (EltR sh') -> SmartExp (EltR sh))
 -> SmartAcc (Array (EltR sh) (EltR a))
 -> PreSmartAcc SmartAcc SmartExp (Array (EltR sh') (EltR a)))
-> ShapeR (EltR sh')
-> SmartExp (EltR sh')
-> (SmartExp (EltR sh') -> SmartExp (EltR sh))
-> SmartAcc (Array (EltR sh) (EltR a))
-> PreSmartAcc SmartAcc SmartExp (Array (EltR sh') (EltR a))
forall a b. (a -> b) -> a -> b
$ Shape sh' => ShapeR (EltR sh')
forall sh. Shape sh => ShapeR (EltR sh)
shapeR @sh')

-- Stencil operations
-- ------------------

-- Common stencil types
--

-- DIM1 stencil type
type Stencil3 a = (Exp a, Exp a, Exp a)
type Stencil5 a = (Exp a, Exp a, Exp a, Exp a, Exp a)
type Stencil7 a = (Exp a, Exp a, Exp a, Exp a, Exp a, Exp a, Exp a)
type Stencil9 a = (Exp a, Exp a, Exp a, Exp a, Exp a, Exp a, Exp a, Exp a, Exp a)

-- DIM2 stencil type
type Stencil3x3 a = (Stencil3 a, Stencil3 a, Stencil3 a)
type Stencil5x3 a = (Stencil5 a, Stencil5 a, Stencil5 a)
type Stencil3x5 a = (Stencil3 a, Stencil3 a, Stencil3 a, Stencil3 a, Stencil3 a)
type Stencil5x5 a = (Stencil5 a, Stencil5 a, Stencil5 a, Stencil5 a, Stencil5 a)

-- DIM3 stencil type
type Stencil3x3x3 a = (Stencil3x3 a, Stencil3x3 a, Stencil3x3 a)
type Stencil5x3x3 a = (Stencil5x3 a, Stencil5x3 a, Stencil5x3 a)
type Stencil3x5x3 a = (Stencil3x5 a, Stencil3x5 a, Stencil3x5 a)
type Stencil3x3x5 a = (Stencil3x3 a, Stencil3x3 a, Stencil3x3 a, Stencil3x3 a, Stencil3x3 a)
type Stencil5x5x3 a = (Stencil5x5 a, Stencil5x5 a, Stencil5x5 a)
type Stencil5x3x5 a = (Stencil5x3 a, Stencil5x3 a, Stencil5x3 a, Stencil5x3 a, Stencil5x3 a)
type Stencil3x5x5 a = (Stencil3x5 a, Stencil3x5 a, Stencil3x5 a, Stencil3x5 a, Stencil3x5 a)
type Stencil5x5x5 a = (Stencil5x5 a, Stencil5x5 a, Stencil5x5 a, Stencil5x5 a, Stencil5x5 a)


-- | Map a stencil over an array. In contrast to 'map', the domain of a stencil
-- function is an entire /neighbourhood/ of each array element. Neighbourhoods
-- are sub-arrays centred around a focal point. They are not necessarily
-- rectangular, but they are symmetric and have an extent of at least three
-- along each axis. Due to the symmetry requirement the extent is necessarily
-- odd. The focal point is the array position that is determined by the stencil.
--
-- For those array positions where the neighbourhood extends past the boundaries
-- of the source array, a boundary condition determines the contents of the
-- out-of-bounds neighbourhood positions.
--
-- Stencil neighbourhoods are specified via nested tuples, where the nesting
-- depth is equal to the dimensionality of the array. For example, a 3x1 stencil
-- for a one-dimensional array:
--
-- > s31 :: Stencil3 a -> Exp a
-- > s31 (l,c,r) = ...
--
-- ...where @c@ is the focal point of the stencil, and @l@ and @r@ represent the
-- elements to the left and right of the focal point, respectively. Similarly,
-- a 3x3 stencil for a two-dimensional array:
--
-- > s33 :: Stencil3x3 a -> Exp a
-- > s33 ((_,t,_)
-- >     ,(l,c,r)
-- >     ,(_,b,_)) = ...
--
-- ...where @c@ is again the focal point and @t@, @b@, @l@ and @r@ are the
-- elements to the top, bottom, left, and right of the focal point, respectively
-- (the diagonal elements have been elided).
--
-- For example, the following computes a 5x5
-- <https://en.wikipedia.org/wiki/Gaussian_blur Gaussian blur> as a separable
-- 2-pass operation.
--
-- > type Stencil5x1 a = (Stencil3 a, Stencil5 a, Stencil3 a)
-- > type Stencil1x5 a = (Stencil3 a, Stencil3 a, Stencil3 a, Stencil3 a, Stencil3 a)
-- >
-- > convolve5x1 :: Num a => [Exp a] -> Stencil5x1 a -> Exp a
-- > convolve5x1 kernel (_, (a,b,c,d,e), _)
-- >   = Prelude.sum $ Prelude.zipWith (*) kernel [a,b,c,d,e]
-- >
-- > convolve1x5 :: Num a => [Exp a] -> Stencil1x5 a -> Exp a
-- > convolve1x5 kernel ((_,a,_), (_,b,_), (_,c,_), (_,d,_), (_,e,_))
-- >   = Prelude.sum $ Prelude.zipWith (*) kernel [a,b,c,d,e]
-- >
-- > gaussian = [0.06136,0.24477,0.38774,0.24477,0.06136]
-- >
-- > blur :: Num a => Acc (Matrix a) -> Acc (Matrix a)
-- > blur = stencil (convolve5x1 gaussian) clamp
-- >      . stencil (convolve1x5 gaussian) clamp
--
-- [/Note:/]
--
-- Since accelerate-1.3.0.0, we allow the source array to fuse into the stencil
-- operation. However, since a stencil computation (typically) requires multiple
-- values from the source array, this means that the work of the fused operation
-- will be duplicated for each element in the stencil pattern.
--
-- For example, suppose we write:
--
-- > blur . map f
--
-- The operation `f` will be fused into each element of the first Gaussian blur
-- kernel, resulting in a stencil equivalent to:
--
-- > f_and_convolve1x5 :: Num a => (Exp a -> Exp b) -> [Exp b] -> Stencil1x5 a -> Exp b
-- > f_and_convolve1x5 f kernel ((_,a,_), (_,b,_), (_,c,_), (_,d,_), (_,e,_))
-- >   = Prelude.sum $ Prelude.zipWith (*) kernel [f a, f b, f c, f d, f e]
--
-- This duplication is often beneficial, however you may choose to instead force
-- the array to be evaluated first, preventing fusion, using the
-- `Data.Array.Accelerate.Prelude.compute` operation. Benchmarking should reveal
-- which approach is best for your application.
--
stencil
    :: forall sh stencil a b.
       (Stencil sh a stencil, Elt b)
    => (stencil -> Exp b)                     -- ^ stencil function
    -> Boundary (Array sh a)                  -- ^ boundary condition
    -> Acc (Array sh a)                       -- ^ source array
    -> Acc (Array sh b)                       -- ^ destination array
stencil :: (stencil -> Exp b)
-> Boundary (Array sh a) -> Acc (Array sh a) -> Acc (Array sh b)
stencil stencil -> Exp b
f (Boundary PreBoundary SmartAcc SmartExp (Array (EltR sh) (EltR e))
b) (Acc SmartAcc (ArraysR (Array sh a))
a)
  = SmartAcc (ArraysR (Array sh b)) -> Acc (Array sh b)
forall a. SmartAcc (ArraysR a) -> Acc a
Acc (SmartAcc (ArraysR (Array sh b)) -> Acc (Array sh b))
-> SmartAcc (ArraysR (Array sh b)) -> Acc (Array sh b)
forall a b. (a -> b) -> a -> b
$ PreSmartAcc SmartAcc SmartExp (Array (EltR sh) (EltR b))
-> SmartAcc (Array (EltR sh) (EltR b))
forall a. PreSmartAcc SmartAcc SmartExp a -> SmartAcc a
SmartAcc (PreSmartAcc SmartAcc SmartExp (Array (EltR sh) (EltR b))
 -> SmartAcc (Array (EltR sh) (EltR b)))
-> PreSmartAcc SmartAcc SmartExp (Array (EltR sh) (EltR b))
-> SmartAcc (Array (EltR sh) (EltR b))
forall a b. (a -> b) -> a -> b
$ StencilR (EltR sh) (EltR a) (StencilR sh stencil)
-> TypeR (EltR b)
-> (SmartExp (StencilR sh stencil) -> SmartExp (EltR b))
-> PreBoundary SmartAcc SmartExp (Array (EltR sh) (EltR a))
-> SmartAcc (Array (EltR sh) (EltR a))
-> PreSmartAcc SmartAcc SmartExp (Array (EltR sh) (EltR b))
forall sh a stencil sh (exp :: * -> *) (acc :: * -> *).
StencilR sh a stencil
-> TypeR sh
-> (SmartExp stencil -> exp sh)
-> PreBoundary acc exp (Array sh a)
-> acc (Array sh a)
-> PreSmartAcc acc exp (Array sh sh)
Stencil
      (Stencil sh a stencil =>
StencilR (EltR sh) (EltR a) (StencilR sh stencil)
forall sh e stencil.
Stencil sh e stencil =>
StencilR (EltR sh) (EltR e) (StencilR sh stencil)
stencilR @sh @a @stencil)
      (Elt b => TypeR (EltR b)
forall a. Elt a => TypeR (EltR a)
eltR @b)
      (Exp b -> SmartExp (EltR b)
forall e. Exp e -> SmartExp (EltR e)
unExp (Exp b -> SmartExp (EltR b))
-> (SmartExp (StencilR sh stencil) -> Exp b)
-> SmartExp (StencilR sh stencil)
-> SmartExp (EltR b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. stencil -> Exp b
f (stencil -> Exp b)
-> (SmartExp (StencilR sh stencil) -> stencil)
-> SmartExp (StencilR sh stencil)
-> Exp b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stencil sh a stencil => SmartExp (StencilR sh stencil) -> stencil
forall sh e stencil.
Stencil sh e stencil =>
SmartExp (StencilR sh stencil) -> stencil
stencilPrj @sh @a @stencil)
      PreBoundary SmartAcc SmartExp (Array (EltR sh) (EltR a))
PreBoundary SmartAcc SmartExp (Array (EltR sh) (EltR e))
b
      SmartAcc (Array (EltR sh) (EltR a))
SmartAcc (ArraysR (Array sh a))
a

-- | Map a binary stencil of an array. The extent of the resulting array is the
-- intersection of the extents of the two source arrays. This is the stencil
-- equivalent of 'zipWith'.
--
stencil2
    :: forall sh stencil1 stencil2 a b c.
       (Stencil sh a stencil1, Stencil sh b stencil2, Elt c)
    => (stencil1 -> stencil2 -> Exp c)        -- ^ binary stencil function
    -> Boundary (Array sh a)                  -- ^ boundary condition #1
    -> Acc (Array sh a)                       -- ^ source array #1
    -> Boundary (Array sh b)                  -- ^ boundary condition #2
    -> Acc (Array sh b)                       -- ^ source array #2
    -> Acc (Array sh c)                       -- ^ destination array
stencil2 :: (stencil1 -> stencil2 -> Exp c)
-> Boundary (Array sh a)
-> Acc (Array sh a)
-> Boundary (Array sh b)
-> Acc (Array sh b)
-> Acc (Array sh c)
stencil2 stencil1 -> stencil2 -> Exp c
f (Boundary PreBoundary SmartAcc SmartExp (Array (EltR sh) (EltR e))
b1) (Acc SmartAcc (ArraysR (Array sh a))
a1) (Boundary PreBoundary SmartAcc SmartExp (Array (EltR sh) (EltR e))
b2) (Acc SmartAcc (ArraysR (Array sh b))
a2)
  = SmartAcc (ArraysR (Array sh c)) -> Acc (Array sh c)
forall a. SmartAcc (ArraysR a) -> Acc a
Acc (SmartAcc (ArraysR (Array sh c)) -> Acc (Array sh c))
-> SmartAcc (ArraysR (Array sh c)) -> Acc (Array sh c)
forall a b. (a -> b) -> a -> b
$ PreSmartAcc SmartAcc SmartExp (Array (EltR sh) (EltR c))
-> SmartAcc (Array (EltR sh) (EltR c))
forall a. PreSmartAcc SmartAcc SmartExp a -> SmartAcc a
SmartAcc (PreSmartAcc SmartAcc SmartExp (Array (EltR sh) (EltR c))
 -> SmartAcc (Array (EltR sh) (EltR c)))
-> PreSmartAcc SmartAcc SmartExp (Array (EltR sh) (EltR c))
-> SmartAcc (Array (EltR sh) (EltR c))
forall a b. (a -> b) -> a -> b
$ StencilR (EltR sh) (EltR a) (StencilR sh stencil1)
-> StencilR (EltR sh) (EltR b) (StencilR sh stencil2)
-> TypeR (EltR c)
-> (SmartExp (StencilR sh stencil1)
    -> SmartExp (StencilR sh stencil2) -> SmartExp (EltR c))
-> PreBoundary SmartAcc SmartExp (Array (EltR sh) (EltR a))
-> SmartAcc (Array (EltR sh) (EltR a))
-> PreBoundary SmartAcc SmartExp (Array (EltR sh) (EltR b))
-> SmartAcc (Array (EltR sh) (EltR b))
-> PreSmartAcc SmartAcc SmartExp (Array (EltR sh) (EltR c))
forall sh a stencil1 b stencil2 c (exp :: * -> *) (acc :: * -> *).
StencilR sh a stencil1
-> StencilR sh b stencil2
-> TypeR c
-> (SmartExp stencil1 -> SmartExp stencil2 -> exp c)
-> PreBoundary acc exp (Array sh a)
-> acc (Array sh a)
-> PreBoundary acc exp (Array sh b)
-> acc (Array sh b)
-> PreSmartAcc acc exp (Array sh c)
Stencil2
      (Stencil sh a stencil1 =>
StencilR (EltR sh) (EltR a) (StencilR sh stencil1)
forall sh e stencil.
Stencil sh e stencil =>
StencilR (EltR sh) (EltR e) (StencilR sh stencil)
stencilR @sh @a @stencil1)
      (Stencil sh b stencil2 =>
StencilR (EltR sh) (EltR b) (StencilR sh stencil2)
forall sh e stencil.
Stencil sh e stencil =>
StencilR (EltR sh) (EltR e) (StencilR sh stencil)
stencilR @sh @b @stencil2)
      (Elt c => TypeR (EltR c)
forall a. Elt a => TypeR (EltR a)
eltR @c)
      (\SmartExp (StencilR sh stencil1)
x SmartExp (StencilR sh stencil2)
y -> Exp c -> SmartExp (EltR c)
forall e. Exp e -> SmartExp (EltR e)
unExp (Exp c -> SmartExp (EltR c)) -> Exp c -> SmartExp (EltR c)
forall a b. (a -> b) -> a -> b
$ stencil1 -> stencil2 -> Exp c
f (SmartExp (StencilR sh stencil1) -> stencil1
forall sh e stencil.
Stencil sh e stencil =>
SmartExp (StencilR sh stencil) -> stencil
stencilPrj @sh @a @stencil1 SmartExp (StencilR sh stencil1)
x) (SmartExp (StencilR sh stencil2) -> stencil2
forall sh e stencil.
Stencil sh e stencil =>
SmartExp (StencilR sh stencil) -> stencil
stencilPrj @sh @b @stencil2 SmartExp (StencilR sh stencil2)
y))
      PreBoundary SmartAcc SmartExp (Array (EltR sh) (EltR a))
PreBoundary SmartAcc SmartExp (Array (EltR sh) (EltR e))
b1
      SmartAcc (Array (EltR sh) (EltR a))
SmartAcc (ArraysR (Array sh a))
a1
      PreBoundary SmartAcc SmartExp (Array (EltR sh) (EltR b))
PreBoundary SmartAcc SmartExp (Array (EltR sh) (EltR e))
b2
      SmartAcc (Array (EltR sh) (EltR b))
SmartAcc (ArraysR (Array sh b))
a2

-- | Boundary condition where elements of the stencil which would be
-- out-of-bounds are instead clamped to the edges of the array.
--
-- In the following 3x3 stencil, the out-of-bounds element @b@ will instead
-- return the value at position @c@:
--
-- >   +------------+
-- >   |a           |
-- >  b|cd          |
-- >   |e           |
-- >   +------------+
--
clamp :: Boundary (Array sh e)
clamp :: Boundary (Array sh e)
clamp = PreBoundary SmartAcc SmartExp (Array (EltR sh) (EltR e))
-> Boundary (Array sh e)
forall sh e.
PreBoundary SmartAcc SmartExp (Array (EltR sh) (EltR e))
-> Boundary (Array sh e)
Boundary PreBoundary SmartAcc SmartExp (Array (EltR sh) (EltR e))
forall (acc :: * -> *) (exp :: * -> *) t. PreBoundary acc exp t
Clamp

-- | Stencil boundary condition where coordinates beyond the array extent are
-- instead mirrored
--
-- In the following 5x3 stencil, the out-of-bounds element @c@ will instead
-- return the value at position @d@, and similarly the element at @b@ will
-- return the value at @e@:
--
-- >   +------------+
-- >   |a           |
-- > bc|def         |
-- >   |g           |
-- >   +------------+
--
mirror :: Boundary (Array sh e)
mirror :: Boundary (Array sh e)
mirror = PreBoundary SmartAcc SmartExp (Array (EltR sh) (EltR e))
-> Boundary (Array sh e)
forall sh e.
PreBoundary SmartAcc SmartExp (Array (EltR sh) (EltR e))
-> Boundary (Array sh e)
Boundary PreBoundary SmartAcc SmartExp (Array (EltR sh) (EltR e))
forall (acc :: * -> *) (exp :: * -> *) t. PreBoundary acc exp t
Mirror

-- | Stencil boundary condition where coordinates beyond the array extent
-- instead wrap around the array (circular boundary conditions).
--
-- In the following 3x3 stencil, the out of bounds elements will be read as in
-- the pattern on the right.
--
-- >  a bc
-- >   +------------+      +------------+
-- >  d|ef          |      |ef         d|
-- >  g|hi          |  ->  |hi         g|
-- >   |            |      |bc         a|
-- >   +------------+      +------------+
--
wrap :: Boundary (Array sh e)
wrap :: Boundary (Array sh e)
wrap = PreBoundary SmartAcc SmartExp (Array (EltR sh) (EltR e))
-> Boundary (Array sh e)
forall sh e.
PreBoundary SmartAcc SmartExp (Array (EltR sh) (EltR e))
-> Boundary (Array sh e)
Boundary PreBoundary SmartAcc SmartExp (Array (EltR sh) (EltR e))
forall (acc :: * -> *) (exp :: * -> *) t. PreBoundary acc exp t
Wrap

-- | Stencil boundary condition where the given function is applied to any
-- outlying coordinates.
--
-- The function is passed the out-of-bounds index, so you can use it to specify
-- different boundary conditions at each side. For example, the following would
-- clamp out-of-bounds elements in the y-direction to zero, while having
-- circular boundary conditions in the x-direction.
--
-- > ring :: Acc (Matrix Float) -> Acc (Matrix Float)
-- > ring xs = stencil f boundary xs
-- >   where
-- >     boundary :: Boundary (Matrix Float)
-- >     boundary = function $ \(unlift -> Z :. y :. x) ->
-- >       if y < 0 || y >= height
-- >         then 0
-- >         else if x < 0
-- >                then xs ! index2 y (width+x)
-- >                else xs ! index2 y (x-width)
-- >
-- >     f :: Stencil3x3 Float -> Exp Float
-- >     f = ...
-- >
-- >     Z :. height :. width = unlift (shape xs)
--
function
    :: forall sh e. (Shape sh, Elt e)
    => (Exp sh -> Exp e)
    -> Boundary (Array sh e)
function :: (Exp sh -> Exp e) -> Boundary (Array sh e)
function Exp sh -> Exp e
f = PreBoundary SmartAcc SmartExp (Array (EltR sh) (EltR e))
-> Boundary (Array sh e)
forall sh e.
PreBoundary SmartAcc SmartExp (Array (EltR sh) (EltR e))
-> Boundary (Array sh e)
Boundary (PreBoundary SmartAcc SmartExp (Array (EltR sh) (EltR e))
 -> Boundary (Array sh e))
-> PreBoundary SmartAcc SmartExp (Array (EltR sh) (EltR e))
-> Boundary (Array sh e)
forall a b. (a -> b) -> a -> b
$ (SmartExp (EltR sh) -> SmartExp (EltR e))
-> PreBoundary SmartAcc SmartExp (Array (EltR sh) (EltR e))
forall sh (exp :: * -> *) e (acc :: * -> *).
(SmartExp sh -> exp e) -> PreBoundary acc exp (Array sh e)
Function (SmartExp (EltR sh) -> SmartExp (EltR e)
f')
  where
    f' :: SmartExp (EltR sh) -> SmartExp (EltR e)
    f' :: SmartExp (EltR sh) -> SmartExp (EltR e)
f' = Exp e -> SmartExp (EltR e)
forall e. Exp e -> SmartExp (EltR e)
unExp (Exp e -> SmartExp (EltR e))
-> (SmartExp (EltR sh) -> Exp e)
-> SmartExp (EltR sh)
-> SmartExp (EltR e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp sh -> Exp e
f (Exp sh -> Exp e)
-> (SmartExp (EltR sh) -> Exp sh) -> SmartExp (EltR sh) -> Exp e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SmartExp (EltR sh) -> Exp sh
forall t. SmartExp (EltR t) -> Exp t
Exp


{--
-- Sequence operations
-- ------------------

-- Common sequence types
--

streamIn :: Arrays a
         => [a]
         -> Seq [a]
streamIn arrs = Seq (StreamIn arrs)

-- | Convert the given array to a sequence by dividing the array up into subarrays.
-- The first argument captures how to the division should be performed. The
-- presence of `All` in the division descriptor indicates that elements in the
-- corresponding dimension should be retained in the subarrays, whereas `Split`
-- indicates that the input array should divided along this dimension.
--
toSeq :: (Division slsix, Elt a)
      => slsix
      -> Acc (Array (FullShape (DivisionSlice slsix)) a)
      -> Seq [Array (SliceShape (DivisionSlice slsix)) a]
toSeq spec acc = Seq (ToSeq spec acc)

-- | Apply the given array function element-wise to the given sequence.
--
mapSeq :: (Arrays a, Arrays b)
       => (Acc a -> Acc b)
       -> Seq [a]
       -> Seq [b]
mapSeq = Seq $$ MapSeq

-- | Apply the given binary function element-wise to the two sequences.  The length of the resulting
-- sequence is the minumum of the lengths of the two source sequences.
--
zipWithSeq :: (Arrays a, Arrays b, Arrays c)
           => (Acc a -> Acc b -> Acc c)
           -> Seq [a]
           -> Seq [b]
           -> Seq [c]
zipWithSeq = Seq $$$ ZipWithSeq

-- | scanSeq (+) a0 x seq. Scan a sequence x by combining each
-- element using the given binary operation (+). (+) must be
-- associative:
--
--   Forall a b c. (a + b) + c = a + (b + c),
--
-- and a0 must be the identity element for (+):
--
--   Forall a. a0 + a = a = a + a0.
--
scanSeq :: Elt a
        => (Exp a -> Exp a -> Exp a)
        -> Exp a
        -> Seq [Scalar a]
        -> Seq [Scalar a]
scanSeq = Seq $$$ ScanSeq

-- | foldSeq (+) a0 x seq. Fold a sequence x by combining each
-- element using the given binary operation (+). (+) must be
-- associative:
--
--   Forall a b c. (a + b) + c = a + (b + c),
--
-- and a0 must be the identity element for (+):
--
--   Forall a. a0 + a = a = a + a0.
--
foldSeq :: Elt a
        => (Exp a -> Exp a -> Exp a)
        -> Exp a
        -> Seq [Scalar a]
        -> Seq (Scalar a)
foldSeq = Seq $$$ FoldSeq

-- | foldSeqFlatten f a0 x seq. A specialized version of
-- FoldSeqAct where reduction with the companion operator
-- corresponds to flattening. f must be semi-associative, with vecotor
-- append (++) as the companion operator:
--
--   Forall b sh1 a1 sh2 a2.
--     f (f b sh1 a1) sh2 a2 = f b (sh1 ++ sh2) (a1 ++ a2).
--
-- It is common to ignore the shape vectors, yielding the usual
-- semi-associativity law:
--
--   f b a _ = b + a,
--
-- for some (+) satisfying:
--
--   Forall b a1 a2. (b + a1) + a2 = b + (a1 ++ a2).
--
foldSeqFlatten :: (Arrays a, Shape jx, Elt b)
               => (Acc a -> Acc (Vector jx) -> Acc (Vector b) -> Acc a)
               -> Acc a
               -> Seq [Array jx b]
               -> Seq a
foldSeqFlatten = Seq $$$ FoldSeqFlatten

collect :: Arrays arrs => Seq arrs -> Acc arrs
collect = Acc . Collect
--}

-- Foreign function calling
-- ------------------------

-- | Call a foreign array function.
--
-- The form the first argument takes is dependent on the backend being targeted.
-- Note that the foreign function only has access to the input array(s) passed
-- in as its argument.
--
-- In case the operation is being executed on a backend which does not support
-- this foreign implementation, the fallback implementation is used instead,
-- which itself could be a foreign implementation for a (presumably) different
-- backend, or an implementation in pure Accelerate. In this way, multiple
-- foreign implementations can be supplied, and will be tested for suitability
-- against the target backend in sequence.
--
-- For an example see the <https://hackage.haskell.org/package/accelerate-fft accelerate-fft> package.
--
foreignAcc
    :: forall as bs asm. (Arrays as, Arrays bs, Foreign asm)
    => asm (ArraysR as -> ArraysR bs)
    -> (Acc as -> Acc bs)
    -> Acc as
    -> Acc bs
foreignAcc :: asm (ArraysR as -> ArraysR bs)
-> (Acc as -> Acc bs) -> Acc as -> Acc bs
foreignAcc asm (ArraysR as -> ArraysR bs)
asm Acc as -> Acc bs
f (Acc SmartAcc (ArraysR as)
as) = SmartAcc (ArraysR bs) -> Acc bs
forall a. SmartAcc (ArraysR a) -> Acc a
Acc (SmartAcc (ArraysR bs) -> Acc bs)
-> SmartAcc (ArraysR bs) -> Acc bs
forall a b. (a -> b) -> a -> b
$ PreSmartAcc SmartAcc SmartExp (ArraysR bs) -> SmartAcc (ArraysR bs)
forall a. PreSmartAcc SmartAcc SmartExp a -> SmartAcc a
SmartAcc (PreSmartAcc SmartAcc SmartExp (ArraysR bs)
 -> SmartAcc (ArraysR bs))
-> PreSmartAcc SmartAcc SmartExp (ArraysR bs)
-> SmartAcc (ArraysR bs)
forall a b. (a -> b) -> a -> b
$ ArraysR (ArraysR bs)
-> asm (ArraysR as -> ArraysR bs)
-> (SmartAcc (ArraysR as) -> SmartAcc (ArraysR bs))
-> SmartAcc (ArraysR as)
-> PreSmartAcc SmartAcc SmartExp (ArraysR bs)
forall (asm :: * -> *) bs as (acc :: * -> *) (exp :: * -> *).
Foreign asm =>
ArraysR bs
-> asm (as -> bs)
-> (SmartAcc as -> SmartAcc bs)
-> acc as
-> PreSmartAcc acc exp bs
Aforeign (Arrays bs => ArraysR (ArraysR bs)
forall a. Arrays a => ArraysR (ArraysR a)
arraysR @bs) asm (ArraysR as -> ArraysR bs)
asm ((Acc as -> Acc bs)
-> SmartAcc (ArraysR as) -> SmartAcc (ArraysR bs)
forall a b.
(Arrays a, Arrays b) =>
(Acc a -> Acc b) -> SmartAcc (ArraysR a) -> SmartAcc (ArraysR b)
unAccFunction Acc as -> Acc bs
f) SmartAcc (ArraysR as)
as

-- | Call a foreign scalar expression.
--
-- The form of the first argument is dependent on the backend being targeted.
-- Note that the foreign function only has access to the input element(s) passed
-- in as its first argument.
--
-- As with 'foreignAcc', the fallback implementation itself may be a (sequence
-- of) foreign implementation(s) for a different backend(s), or implemented
-- purely in Accelerate.
--
foreignExp
    :: forall x y asm. (Elt x, Elt y, Foreign asm)
    => asm (EltR x -> EltR y)
    -> (Exp x -> Exp y)
    -> Exp x
    -> Exp y
foreignExp :: asm (EltR x -> EltR y) -> (Exp x -> Exp y) -> Exp x -> Exp y
foreignExp asm (EltR x -> EltR y)
asm Exp x -> Exp y
f (Exp SmartExp (EltR x)
x) = PreSmartExp SmartAcc SmartExp (EltR y) -> Exp y
forall t. PreSmartExp SmartAcc SmartExp (EltR t) -> Exp t
mkExp (PreSmartExp SmartAcc SmartExp (EltR y) -> Exp y)
-> PreSmartExp SmartAcc SmartExp (EltR y) -> Exp y
forall a b. (a -> b) -> a -> b
$ TypeR (EltR y)
-> asm (EltR x -> EltR y)
-> (SmartExp (EltR x) -> SmartExp (EltR y))
-> SmartExp (EltR x)
-> PreSmartExp SmartAcc SmartExp (EltR y)
forall (asm :: * -> *) y x (exp :: * -> *) (acc :: * -> *).
Foreign asm =>
TypeR y
-> asm (x -> y)
-> (SmartExp x -> SmartExp y)
-> exp x
-> PreSmartExp acc exp y
Foreign (Elt y => TypeR (EltR y)
forall a. Elt a => TypeR (EltR a)
eltR @y) asm (EltR x -> EltR y)
asm ((Exp x -> Exp y) -> SmartExp (EltR x) -> SmartExp (EltR y)
forall a b.
(Elt a, Elt b) =>
(Exp a -> Exp b) -> SmartExp (EltR a) -> SmartExp (EltR b)
unExpFunction Exp x -> Exp y
f) SmartExp (EltR x)
x


-- Composition of array computations
-- ---------------------------------

-- | Pipelining of two array computations. The first argument will be fully
-- evaluated before being passed to the second computation. This can be used to
-- prevent the argument being fused into the function, for example.
--
-- Denotationally, we have
--
-- > (acc1 >-> acc2) arrs = let tmp = acc1 arrs
-- >                        in  tmp `seq` acc2 tmp
--
-- For an example use of this operation see the 'Data.Array.Accelerate.compute'
-- function.
--
infixl 1 >->
(>->) :: forall a b c. (Arrays a, Arrays b, Arrays c) => (Acc a -> Acc b) -> (Acc b -> Acc c) -> (Acc a -> Acc c)
>-> :: (Acc a -> Acc b) -> (Acc b -> Acc c) -> Acc a -> Acc c
(>->) = SmartAcc (ArraysR c) -> Acc c
forall a. SmartAcc (ArraysR a) -> Acc a
Acc (SmartAcc (ArraysR c) -> Acc c)
-> ((Acc a -> Acc b)
    -> (Acc b -> Acc c) -> Acc a -> SmartAcc (ArraysR c))
-> (Acc a -> Acc b)
-> (Acc b -> Acc c)
-> Acc a
-> Acc c
forall b a c d e.
(b -> a) -> (c -> d -> e -> b) -> c -> d -> e -> a
$$$ FromApplyAcc
  ((Acc a -> Acc b)
   -> (Acc b -> Acc c) -> Acc a -> SmartAcc (ArraysR c))
-> (Acc a -> Acc b)
-> (Acc b -> Acc c)
-> Acc a
-> SmartAcc (ArraysR c)
forall a. ApplyAcc a => FromApplyAcc a -> a
applyAcc (FromApplyAcc
   ((Acc a -> Acc b)
    -> (Acc b -> Acc c) -> Acc a -> SmartAcc (ArraysR c))
 -> (Acc a -> Acc b)
 -> (Acc b -> Acc c)
 -> Acc a
 -> SmartAcc (ArraysR c))
-> FromApplyAcc
     ((Acc a -> Acc b)
      -> (Acc b -> Acc c) -> Acc a -> SmartAcc (ArraysR c))
-> (Acc a -> Acc b)
-> (Acc b -> Acc c)
-> Acc a
-> SmartAcc (ArraysR c)
forall a b. (a -> b) -> a -> b
$ ArraysR (ArraysR a)
-> ArraysR (ArraysR b)
-> ArraysR (ArraysR c)
-> (SmartAcc (ArraysR a) -> SmartAcc (ArraysR b))
-> (SmartAcc (ArraysR b) -> SmartAcc (ArraysR c))
-> SmartAcc (ArraysR a)
-> PreSmartAcc SmartAcc SmartExp (ArraysR c)
forall as bs cs (acc :: * -> *) (exp :: * -> *).
ArraysR as
-> ArraysR bs
-> ArraysR cs
-> (SmartAcc as -> acc bs)
-> (SmartAcc bs -> acc cs)
-> acc as
-> PreSmartAcc acc exp cs
Pipe (Arrays a => ArraysR (ArraysR a)
forall a. Arrays a => ArraysR (ArraysR a)
arraysR @a) (Arrays b => ArraysR (ArraysR b)
forall a. Arrays a => ArraysR (ArraysR a)
arraysR @b) (Arrays c => ArraysR (ArraysR c)
forall a. Arrays a => ArraysR (ArraysR a)
arraysR @c)


-- Flow control constructs
-- -----------------------

-- | An array-level if-then-else construct.
--
-- Enabling the @RebindableSyntax@ extension will allow you to use the standard
-- if-then-else syntax instead.
--
acond :: Arrays a
      => Exp Bool               -- ^ if-condition
      -> Acc a                  -- ^ then-array
      -> Acc a                  -- ^ else-array
      -> Acc a
acond :: Exp Bool -> Acc a -> Acc a -> Acc a
acond (Exp SmartExp (EltR Bool)
p) = SmartAcc (ArraysR a) -> Acc a
forall a. SmartAcc (ArraysR a) -> Acc a
Acc (SmartAcc (ArraysR a) -> Acc a)
-> (Acc a -> Acc a -> SmartAcc (ArraysR a))
-> Acc a
-> Acc a
-> Acc a
forall b a c d. (b -> a) -> (c -> d -> b) -> c -> d -> a
$$ FromApplyAcc (Acc a -> Acc a -> SmartAcc (ArraysR a))
-> Acc a -> Acc a -> SmartAcc (ArraysR a)
forall a. ApplyAcc a => FromApplyAcc a -> a
applyAcc (FromApplyAcc (Acc a -> Acc a -> SmartAcc (ArraysR a))
 -> Acc a -> Acc a -> SmartAcc (ArraysR a))
-> FromApplyAcc (Acc a -> Acc a -> SmartAcc (ArraysR a))
-> Acc a
-> Acc a
-> SmartAcc (ArraysR a)
forall a b. (a -> b) -> a -> b
$ SmartExp PrimBool
-> SmartAcc (ArraysR a)
-> SmartAcc (ArraysR a)
-> PreSmartAcc SmartAcc SmartExp (ArraysR a)
forall (exp :: * -> *) (acc :: * -> *) as.
exp PrimBool -> acc as -> acc as -> PreSmartAcc acc exp as
Acond (SmartExp (PrimBool, ()) -> SmartExp PrimBool
forall a b. Coerce a b => SmartExp a -> SmartExp b
mkCoerce' SmartExp (PrimBool, ())
SmartExp (EltR Bool)
p)

-- | An array-level 'while' construct. Continue to apply the given function,
-- starting with the initial value, until the test function evaluates to
-- 'False'.
--
awhile :: forall a. Arrays a
       => (Acc a -> Acc (Scalar Bool))    -- ^ keep evaluating while this returns 'True'
       -> (Acc a -> Acc a)                -- ^ function to apply
       -> Acc a                           -- ^ initial value
       -> Acc a
awhile :: (Acc a -> Acc (Scalar Bool)) -> (Acc a -> Acc a) -> Acc a -> Acc a
awhile Acc a -> Acc (Scalar Bool)
f = SmartAcc (ArraysR a) -> Acc a
forall a. SmartAcc (ArraysR a) -> Acc a
Acc (SmartAcc (ArraysR a) -> Acc a)
-> ((Acc a -> Acc a) -> Acc a -> SmartAcc (ArraysR a))
-> (Acc a -> Acc a)
-> Acc a
-> Acc a
forall b a c d. (b -> a) -> (c -> d -> b) -> c -> d -> a
$$ FromApplyAcc ((Acc a -> Acc a) -> Acc a -> SmartAcc (ArraysR a))
-> (Acc a -> Acc a) -> Acc a -> SmartAcc (ArraysR a)
forall a. ApplyAcc a => FromApplyAcc a -> a
applyAcc (FromApplyAcc ((Acc a -> Acc a) -> Acc a -> SmartAcc (ArraysR a))
 -> (Acc a -> Acc a) -> Acc a -> SmartAcc (ArraysR a))
-> FromApplyAcc ((Acc a -> Acc a) -> Acc a -> SmartAcc (ArraysR a))
-> (Acc a -> Acc a)
-> Acc a
-> SmartAcc (ArraysR a)
forall a b. (a -> b) -> a -> b
$ ArraysR (ArraysR a)
-> (SmartAcc (ArraysR a) -> SmartAcc (Scalar PrimBool))
-> (SmartAcc (ArraysR a) -> SmartAcc (ArraysR a))
-> SmartAcc (ArraysR a)
-> PreSmartAcc SmartAcc SmartExp (ArraysR a)
forall arrs (acc :: * -> *) (exp :: * -> *).
ArraysR arrs
-> (SmartAcc arrs -> acc (Scalar PrimBool))
-> (SmartAcc arrs -> acc arrs)
-> acc arrs
-> PreSmartAcc acc exp arrs
Awhile (Arrays a => ArraysR (ArraysR a)
forall a. Arrays a => ArraysR (ArraysR a)
arraysR @a) ((Acc a -> Acc (Scalar PrimBool))
-> SmartAcc (ArraysR a) -> SmartAcc (ArraysR (Scalar PrimBool))
forall a b.
(Arrays a, Arrays b) =>
(Acc a -> Acc b) -> SmartAcc (ArraysR a) -> SmartAcc (ArraysR b)
unAccFunction Acc a -> Acc (Scalar PrimBool)
g)
  where
    -- FIXME: This should be a no-op!
    g :: Acc a -> Acc (Scalar PrimBool)
    g :: Acc a -> Acc (Scalar PrimBool)
g = (Exp Bool -> Exp PrimBool)
-> Acc (Scalar Bool) -> Acc (Scalar PrimBool)
forall sh a b.
(Shape sh, Elt a, Elt b) =>
(Exp a -> Exp b) -> Acc (Array sh a) -> Acc (Array sh b)
map Exp Bool -> Exp PrimBool
forall a b. Coerce (EltR a) (EltR b) => Exp a -> Exp b
mkCoerce (Acc (Scalar Bool) -> Acc (Scalar PrimBool))
-> (Acc a -> Acc (Scalar Bool)) -> Acc a -> Acc (Scalar PrimBool)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Acc a -> Acc (Scalar Bool)
f


-- Shapes and indices
-- ------------------

-- | Map a multi-dimensional index into a linear, row-major representation of an
-- array.
--
toIndex
    :: forall sh. Shape sh
    => Exp sh                     -- ^ extent of the array
    -> Exp sh                     -- ^ index to remap
    -> Exp Int
toIndex :: Exp sh -> Exp sh -> Exp Int
toIndex (Exp SmartExp (EltR sh)
sh) (Exp SmartExp (EltR sh)
ix) = PreSmartExp SmartAcc SmartExp (EltR Int) -> Exp Int
forall t. PreSmartExp SmartAcc SmartExp (EltR t) -> Exp t
mkExp (PreSmartExp SmartAcc SmartExp (EltR Int) -> Exp Int)
-> PreSmartExp SmartAcc SmartExp (EltR Int) -> Exp Int
forall a b. (a -> b) -> a -> b
$ ShapeR (EltR sh)
-> SmartExp (EltR sh)
-> SmartExp (EltR sh)
-> PreSmartExp SmartAcc SmartExp Int
forall sh (exp :: * -> *) (acc :: * -> *).
ShapeR sh -> exp sh -> exp sh -> PreSmartExp acc exp Int
ToIndex (Shape sh => ShapeR (EltR sh)
forall sh. Shape sh => ShapeR (EltR sh)
shapeR @sh) SmartExp (EltR sh)
sh SmartExp (EltR sh)
ix

-- | Inverse of 'toIndex'
--
fromIndex :: forall sh. Shape sh => Exp sh -> Exp Int -> Exp sh
fromIndex :: Exp sh -> Exp Int -> Exp sh
fromIndex (Exp SmartExp (EltR sh)
sh) (Exp SmartExp (EltR Int)
e) = PreSmartExp SmartAcc SmartExp (EltR sh) -> Exp sh
forall t. PreSmartExp SmartAcc SmartExp (EltR t) -> Exp t
mkExp (PreSmartExp SmartAcc SmartExp (EltR sh) -> Exp sh)
-> PreSmartExp SmartAcc SmartExp (EltR sh) -> Exp sh
forall a b. (a -> b) -> a -> b
$ ShapeR (EltR sh)
-> SmartExp (EltR sh)
-> SmartExp Int
-> PreSmartExp SmartAcc SmartExp (EltR sh)
forall sh (exp :: * -> *) (acc :: * -> *).
ShapeR sh -> exp sh -> exp Int -> PreSmartExp acc exp sh
FromIndex (Shape sh => ShapeR (EltR sh)
forall sh. Shape sh => ShapeR (EltR sh)
shapeR @sh) SmartExp (EltR sh)
sh SmartExp Int
SmartExp (EltR Int)
e

-- | Intersection of two shapes
--
intersect :: forall sh. Shape sh => Exp sh -> Exp sh -> Exp sh
intersect :: Exp sh -> Exp sh -> Exp sh
intersect (Exp SmartExp (EltR sh)
shx) (Exp SmartExp (EltR sh)
shy) = SmartExp (EltR sh) -> Exp sh
forall t. SmartExp (EltR t) -> Exp t
Exp (SmartExp (EltR sh) -> Exp sh) -> SmartExp (EltR sh) -> Exp sh
forall a b. (a -> b) -> a -> b
$ ShapeR (EltR sh)
-> SmartExp (EltR sh) -> SmartExp (EltR sh) -> SmartExp (EltR sh)
forall t. ShapeR t -> SmartExp t -> SmartExp t -> SmartExp t
intersect' (Shape sh => ShapeR (EltR sh)
forall sh. Shape sh => ShapeR (EltR sh)
shapeR @sh) SmartExp (EltR sh)
shx SmartExp (EltR sh)
shy
  where
    intersect' :: ShapeR t -> SmartExp t -> SmartExp t -> SmartExp t
    intersect' :: ShapeR t -> SmartExp t -> SmartExp t -> SmartExp t
intersect' ShapeR t
ShapeRz SmartExp t
_ SmartExp t
_ = PreSmartExp SmartAcc SmartExp () -> SmartExp ()
forall t. PreSmartExp SmartAcc SmartExp t -> SmartExp t
SmartExp PreSmartExp SmartAcc SmartExp ()
forall (acc :: * -> *) (exp :: * -> *). PreSmartExp acc exp ()
Nil
    intersect' (ShapeRsnoc ShapeR sh
shR) (SmartExp t -> (SmartExp sh, SmartExp Int)
forall a b. SmartExp (a, b) -> (SmartExp a, SmartExp b)
unPair -> (SmartExp sh
xs, SmartExp Int
x)) (SmartExp t -> (SmartExp sh, SmartExp Int)
forall a b. SmartExp (a, b) -> (SmartExp a, SmartExp b)
unPair -> (SmartExp sh
ys, SmartExp Int
y))
      = PreSmartExp SmartAcc SmartExp (sh, Int) -> SmartExp (sh, Int)
forall t. PreSmartExp SmartAcc SmartExp t -> SmartExp t
SmartExp
      (PreSmartExp SmartAcc SmartExp (sh, Int) -> SmartExp (sh, Int))
-> PreSmartExp SmartAcc SmartExp (sh, Int) -> SmartExp (sh, Int)
forall a b. (a -> b) -> a -> b
$ ShapeR sh -> SmartExp sh -> SmartExp sh -> SmartExp sh
forall t. ShapeR t -> SmartExp t -> SmartExp t -> SmartExp t
intersect' ShapeR sh
shR SmartExp sh
xs SmartExp sh
ys SmartExp sh
-> SmartExp Int -> PreSmartExp SmartAcc SmartExp (sh, Int)
forall (exp :: * -> *) t1 t2 (acc :: * -> *).
exp t1 -> exp t2 -> PreSmartExp acc exp (t1, t2)
`Pair`
        PreSmartExp SmartAcc SmartExp Int -> SmartExp Int
forall t. PreSmartExp SmartAcc SmartExp t -> SmartExp t
SmartExp (PrimFun ((Int, Int) -> Int)
-> SmartExp (Int, Int) -> PreSmartExp SmartAcc SmartExp Int
forall a r (exp :: * -> *) (acc :: * -> *).
PrimFun (a -> r) -> exp a -> PreSmartExp acc exp r
PrimApp (SingleType Int -> PrimFun ((Int, Int) -> Int)
forall a. SingleType a -> PrimFun ((a, a) -> a)
PrimMin SingleType Int
forall a. IsSingle a => SingleType a
singleType) (SmartExp (Int, Int) -> PreSmartExp SmartAcc SmartExp Int)
-> SmartExp (Int, Int) -> PreSmartExp SmartAcc SmartExp Int
forall a b. (a -> b) -> a -> b
$ PreSmartExp SmartAcc SmartExp (Int, Int) -> SmartExp (Int, Int)
forall t. PreSmartExp SmartAcc SmartExp t -> SmartExp t
SmartExp (PreSmartExp SmartAcc SmartExp (Int, Int) -> SmartExp (Int, Int))
-> PreSmartExp SmartAcc SmartExp (Int, Int) -> SmartExp (Int, Int)
forall a b. (a -> b) -> a -> b
$ SmartExp Int
-> SmartExp Int -> PreSmartExp SmartAcc SmartExp (Int, Int)
forall (exp :: * -> *) t1 t2 (acc :: * -> *).
exp t1 -> exp t2 -> PreSmartExp acc exp (t1, t2)
Pair SmartExp Int
x SmartExp Int
y)


-- | Union of two shapes
--
union :: forall sh. Shape sh => Exp sh -> Exp sh -> Exp sh
union :: Exp sh -> Exp sh -> Exp sh
union (Exp SmartExp (EltR sh)
shx) (Exp SmartExp (EltR sh)
shy) = SmartExp (EltR sh) -> Exp sh
forall t. SmartExp (EltR t) -> Exp t
Exp (SmartExp (EltR sh) -> Exp sh) -> SmartExp (EltR sh) -> Exp sh
forall a b. (a -> b) -> a -> b
$ ShapeR (EltR sh)
-> SmartExp (EltR sh) -> SmartExp (EltR sh) -> SmartExp (EltR sh)
forall t. ShapeR t -> SmartExp t -> SmartExp t -> SmartExp t
union' (Shape sh => ShapeR (EltR sh)
forall sh. Shape sh => ShapeR (EltR sh)
shapeR @sh) SmartExp (EltR sh)
shx SmartExp (EltR sh)
shy
  where
    union' :: ShapeR t -> SmartExp t -> SmartExp t -> SmartExp t
    union' :: ShapeR t -> SmartExp t -> SmartExp t -> SmartExp t
union' ShapeR t
ShapeRz SmartExp t
_ SmartExp t
_ = PreSmartExp SmartAcc SmartExp () -> SmartExp ()
forall t. PreSmartExp SmartAcc SmartExp t -> SmartExp t
SmartExp PreSmartExp SmartAcc SmartExp ()
forall (acc :: * -> *) (exp :: * -> *). PreSmartExp acc exp ()
Nil
    union' (ShapeRsnoc ShapeR sh
shR) (SmartExp t -> (SmartExp sh, SmartExp Int)
forall a b. SmartExp (a, b) -> (SmartExp a, SmartExp b)
unPair -> (SmartExp sh
xs, SmartExp Int
x)) (SmartExp t -> (SmartExp sh, SmartExp Int)
forall a b. SmartExp (a, b) -> (SmartExp a, SmartExp b)
unPair -> (SmartExp sh
ys, SmartExp Int
y))
      = PreSmartExp SmartAcc SmartExp (sh, Int) -> SmartExp (sh, Int)
forall t. PreSmartExp SmartAcc SmartExp t -> SmartExp t
SmartExp
      (PreSmartExp SmartAcc SmartExp (sh, Int) -> SmartExp (sh, Int))
-> PreSmartExp SmartAcc SmartExp (sh, Int) -> SmartExp (sh, Int)
forall a b. (a -> b) -> a -> b
$ ShapeR sh -> SmartExp sh -> SmartExp sh -> SmartExp sh
forall t. ShapeR t -> SmartExp t -> SmartExp t -> SmartExp t
union' ShapeR sh
shR SmartExp sh
xs SmartExp sh
ys SmartExp sh
-> SmartExp Int -> PreSmartExp SmartAcc SmartExp (sh, Int)
forall (exp :: * -> *) t1 t2 (acc :: * -> *).
exp t1 -> exp t2 -> PreSmartExp acc exp (t1, t2)
`Pair`
        PreSmartExp SmartAcc SmartExp Int -> SmartExp Int
forall t. PreSmartExp SmartAcc SmartExp t -> SmartExp t
SmartExp (PrimFun ((Int, Int) -> Int)
-> SmartExp (Int, Int) -> PreSmartExp SmartAcc SmartExp Int
forall a r (exp :: * -> *) (acc :: * -> *).
PrimFun (a -> r) -> exp a -> PreSmartExp acc exp r
PrimApp (SingleType Int -> PrimFun ((Int, Int) -> Int)
forall a. SingleType a -> PrimFun ((a, a) -> a)
PrimMax SingleType Int
forall a. IsSingle a => SingleType a
singleType) (SmartExp (Int, Int) -> PreSmartExp SmartAcc SmartExp Int)
-> SmartExp (Int, Int) -> PreSmartExp SmartAcc SmartExp Int
forall a b. (a -> b) -> a -> b
$ PreSmartExp SmartAcc SmartExp (Int, Int) -> SmartExp (Int, Int)
forall t. PreSmartExp SmartAcc SmartExp t -> SmartExp t
SmartExp (PreSmartExp SmartAcc SmartExp (Int, Int) -> SmartExp (Int, Int))
-> PreSmartExp SmartAcc SmartExp (Int, Int) -> SmartExp (Int, Int)
forall a b. (a -> b) -> a -> b
$ SmartExp Int
-> SmartExp Int -> PreSmartExp SmartAcc SmartExp (Int, Int)
forall (exp :: * -> *) t1 t2 (acc :: * -> *).
exp t1 -> exp t2 -> PreSmartExp acc exp (t1, t2)
Pair SmartExp Int
x SmartExp Int
y)


-- Flow-control
-- ------------

-- | A scalar-level if-then-else construct.
--
-- Enabling the @RebindableSyntax@ extension will allow you to use the standard
-- if-then-else syntax instead.
--
cond :: Elt t
     => Exp Bool                -- ^ condition
     -> Exp t                   -- ^ then-expression
     -> Exp t                   -- ^ else-expression
     -> Exp t
cond :: Exp Bool -> Exp t -> Exp t -> Exp t
cond (Exp SmartExp (EltR Bool)
c) (Exp SmartExp (EltR t)
x) (Exp SmartExp (EltR t)
y) = PreSmartExp SmartAcc SmartExp (EltR t) -> Exp t
forall t. PreSmartExp SmartAcc SmartExp (EltR t) -> Exp t
mkExp (PreSmartExp SmartAcc SmartExp (EltR t) -> Exp t)
-> PreSmartExp SmartAcc SmartExp (EltR t) -> Exp t
forall a b. (a -> b) -> a -> b
$ SmartExp PrimBool
-> SmartExp (EltR t)
-> SmartExp (EltR t)
-> PreSmartExp SmartAcc SmartExp (EltR t)
forall (exp :: * -> *) t (acc :: * -> *).
exp PrimBool -> exp t -> exp t -> PreSmartExp acc exp t
Cond (SmartExp (PrimBool, ()) -> SmartExp PrimBool
forall a b. Coerce a b => SmartExp a -> SmartExp b
mkCoerce' SmartExp (PrimBool, ())
SmartExp (EltR Bool)
c) SmartExp (EltR t)
x SmartExp (EltR t)
y

-- | While construct. Continue to apply the given function, starting with the
-- initial value, until the test function evaluates to 'False'.
--
while :: forall e. Elt e
      => (Exp e -> Exp Bool)    -- ^ keep evaluating while this returns 'True'
      -> (Exp e -> Exp e)       -- ^ function to apply
      -> Exp e                  -- ^ initial value
      -> Exp e
while :: (Exp e -> Exp Bool) -> (Exp e -> Exp e) -> Exp e -> Exp e
while Exp e -> Exp Bool
c Exp e -> Exp e
f (Exp SmartExp (EltR e)
e) =
  PreSmartExp SmartAcc SmartExp (EltR e) -> Exp e
forall t. PreSmartExp SmartAcc SmartExp (EltR t) -> Exp t
mkExp (PreSmartExp SmartAcc SmartExp (EltR e) -> Exp e)
-> PreSmartExp SmartAcc SmartExp (EltR e) -> Exp e
forall a b. (a -> b) -> a -> b
$ TypeR (EltR e)
-> (SmartExp (EltR e) -> SmartExp PrimBool)
-> (SmartExp (EltR e) -> SmartExp (EltR e))
-> SmartExp (EltR e)
-> PreSmartExp SmartAcc SmartExp (EltR e)
forall t (exp :: * -> *) (acc :: * -> *).
TypeR t
-> (SmartExp t -> exp PrimBool)
-> (SmartExp t -> exp t)
-> exp t
-> PreSmartExp acc exp t
While @(EltR e) (Elt e => TypeR (EltR e)
forall a. Elt a => TypeR (EltR a)
eltR @e)
            (SmartExp (PrimBool, ()) -> SmartExp PrimBool
forall a b. Coerce a b => SmartExp a -> SmartExp b
mkCoerce' (SmartExp (PrimBool, ()) -> SmartExp PrimBool)
-> (SmartExp (EltR e) -> SmartExp (PrimBool, ()))
-> SmartExp (EltR e)
-> SmartExp PrimBool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp Bool -> SmartExp (PrimBool, ())
forall e. Exp e -> SmartExp (EltR e)
unExp (Exp Bool -> SmartExp (PrimBool, ()))
-> (SmartExp (EltR e) -> Exp Bool)
-> SmartExp (EltR e)
-> SmartExp (PrimBool, ())
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp e -> Exp Bool
c (Exp e -> Exp Bool)
-> (SmartExp (EltR e) -> Exp e) -> SmartExp (EltR e) -> Exp Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SmartExp (EltR e) -> Exp e
forall t. SmartExp (EltR t) -> Exp t
Exp)
            (Exp e -> SmartExp (EltR e)
forall e. Exp e -> SmartExp (EltR e)
unExp (Exp e -> SmartExp (EltR e))
-> (SmartExp (EltR e) -> Exp e)
-> SmartExp (EltR e)
-> SmartExp (EltR e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Exp e -> Exp e
f (Exp e -> Exp e)
-> (SmartExp (EltR e) -> Exp e) -> SmartExp (EltR e) -> Exp e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SmartExp (EltR e) -> Exp e
forall t. SmartExp (EltR t) -> Exp t
Exp) SmartExp (EltR e)
e


-- Array operations with a scalar result
-- -------------------------------------

-- | Multidimensional array indexing. Extract the value from an array at the
-- specified zero-based index.
--
-- >>> let mat = fromList (Z:.5:.10) [0..] :: Matrix Int
-- >>> mat
-- Matrix (Z :. 5 :. 10)
--   [  0,  1,  2,  3,  4,  5,  6,  7,  8,  9,
--     10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
--     20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
--     30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
--     40, 41, 42, 43, 44, 45, 46, 47, 48, 49]
--
-- >>> runExp $ use mat ! constant (Z:.1:.2)
-- 12
--
infixl 9 !
(!) :: forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh -> Exp e
Acc SmartAcc (ArraysR (Array sh e))
a ! :: Acc (Array sh e) -> Exp sh -> Exp e
! Exp SmartExp (EltR sh)
ix = PreSmartExp SmartAcc SmartExp (EltR e) -> Exp e
forall t. PreSmartExp SmartAcc SmartExp (EltR t) -> Exp t
mkExp (PreSmartExp SmartAcc SmartExp (EltR e) -> Exp e)
-> PreSmartExp SmartAcc SmartExp (EltR e) -> Exp e
forall a b. (a -> b) -> a -> b
$ TypeR (EltR e)
-> SmartAcc (Array (EltR sh) (EltR e))
-> SmartExp (EltR sh)
-> PreSmartExp SmartAcc SmartExp (EltR e)
forall t (acc :: * -> *) sh (exp :: * -> *).
TypeR t -> acc (Array sh t) -> exp sh -> PreSmartExp acc exp t
Index (Elt e => TypeR (EltR e)
forall a. Elt a => TypeR (EltR a)
eltR @e) SmartAcc (Array (EltR sh) (EltR e))
SmartAcc (ArraysR (Array sh e))
a SmartExp (EltR sh)
ix

-- | Extract the value from an array at the specified linear index.
-- Multidimensional arrays in Accelerate are stored in row-major order with
-- zero-based indexing.
--
-- >>> let mat = fromList (Z:.5:.10) [0..] :: Matrix Int
-- >>> mat
-- Matrix (Z :. 5 :. 10)
--   [  0,  1,  2,  3,  4,  5,  6,  7,  8,  9,
--     10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
--     20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
--     30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
--     40, 41, 42, 43, 44, 45, 46, 47, 48, 49]
--
-- >>> runExp $ use mat !! 12
-- 12
--
infixl 9 !!
(!!) :: forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp Int -> Exp e
Acc SmartAcc (ArraysR (Array sh e))
a !! :: Acc (Array sh e) -> Exp Int -> Exp e
!! Exp SmartExp (EltR Int)
ix = PreSmartExp SmartAcc SmartExp (EltR e) -> Exp e
forall t. PreSmartExp SmartAcc SmartExp (EltR t) -> Exp t
mkExp (PreSmartExp SmartAcc SmartExp (EltR e) -> Exp e)
-> PreSmartExp SmartAcc SmartExp (EltR e) -> Exp e
forall a b. (a -> b) -> a -> b
$ TypeR (EltR e)
-> SmartAcc (Array (EltR sh) (EltR e))
-> SmartExp Int
-> PreSmartExp SmartAcc SmartExp (EltR e)
forall t (acc :: * -> *) sh (exp :: * -> *).
TypeR t -> acc (Array sh t) -> exp Int -> PreSmartExp acc exp t
LinearIndex (Elt e => TypeR (EltR e)
forall a. Elt a => TypeR (EltR a)
eltR @e) SmartAcc (Array (EltR sh) (EltR e))
SmartAcc (ArraysR (Array sh e))
a SmartExp Int
SmartExp (EltR Int)
ix

-- | Extract the shape (extent) of an array.
--
shape :: forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape :: Acc (Array sh e) -> Exp sh
shape = PreSmartExp SmartAcc SmartExp (EltR sh) -> Exp sh
forall t. PreSmartExp SmartAcc SmartExp (EltR t) -> Exp t
mkExp (PreSmartExp SmartAcc SmartExp (EltR sh) -> Exp sh)
-> (Acc (Array sh e) -> PreSmartExp SmartAcc SmartExp (EltR sh))
-> Acc (Array sh e)
-> Exp sh
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShapeR (EltR sh)
-> SmartAcc (Array (EltR sh) (EltR e))
-> PreSmartExp SmartAcc SmartExp (EltR sh)
forall sh (acc :: * -> *) e (exp :: * -> *).
ShapeR sh -> acc (Array sh e) -> PreSmartExp acc exp sh
Shape (Shape sh => ShapeR (EltR sh)
forall sh. Shape sh => ShapeR (EltR sh)
shapeR @sh) (SmartAcc (Array (EltR sh) (EltR e))
 -> PreSmartExp SmartAcc SmartExp (EltR sh))
-> (Acc (Array sh e) -> SmartAcc (Array (EltR sh) (EltR e)))
-> Acc (Array sh e)
-> PreSmartExp SmartAcc SmartExp (EltR sh)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Acc (Array sh e) -> SmartAcc (Array (EltR sh) (EltR e))
forall a. Arrays a => Acc a -> SmartAcc (ArraysR a)
unAcc

-- | The number of elements in the array
--
size :: (Shape sh, Elt e) => Acc (Array sh e) -> Exp Int
size :: Acc (Array sh e) -> Exp Int
size = Exp sh -> Exp Int
forall sh. Shape sh => Exp sh -> Exp Int
shapeSize (Exp sh -> Exp Int)
-> (Acc (Array sh e) -> Exp sh) -> Acc (Array sh e) -> Exp Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Acc (Array sh e) -> Exp sh
forall sh e. (Shape sh, Elt e) => Acc (Array sh e) -> Exp sh
shape

-- | The number of elements that would be held by an array of the given shape.
--
shapeSize :: forall sh. Shape sh => Exp sh -> Exp Int
shapeSize :: Exp sh -> Exp Int
shapeSize (Exp SmartExp (EltR sh)
sh) = PreSmartExp SmartAcc SmartExp (EltR Int) -> Exp Int
forall t. PreSmartExp SmartAcc SmartExp (EltR t) -> Exp t
mkExp (PreSmartExp SmartAcc SmartExp (EltR Int) -> Exp Int)
-> PreSmartExp SmartAcc SmartExp (EltR Int) -> Exp Int
forall a b. (a -> b) -> a -> b
$ ShapeR (EltR sh)
-> SmartExp (EltR sh) -> PreSmartExp SmartAcc SmartExp Int
forall sh (exp :: * -> *) (acc :: * -> *).
ShapeR sh -> exp sh -> PreSmartExp acc exp Int
ShapeSize (Shape sh => ShapeR (EltR sh)
forall sh. Shape sh => ShapeR (EltR sh)
shapeR @sh) SmartExp (EltR sh)
sh


-- Numeric functions
-- -----------------

-- | 'subtract' is the same as @'flip' ('-')@.
--
subtract :: Num a => Exp a -> Exp a -> Exp a
subtract :: Exp a -> Exp a -> Exp a
subtract Exp a
x Exp a
y = Exp a
y Exp a -> Exp a -> Exp a
forall a. Num a => a -> a -> a
- Exp a
x

-- | Determine if a number is even
--
even :: Integral a => Exp a -> Exp Bool
even :: Exp a -> Exp Bool
even Exp a
n = Exp a
n Exp a -> Exp a -> Exp a
forall a. Integral a => a -> a -> a
`rem` Exp a
2 Exp a -> Exp a -> Exp Bool
forall a. Eq a => Exp a -> Exp a -> Exp Bool
== Exp a
0

-- | Determine if a number is odd
--
odd :: Integral a => Exp a -> Exp Bool
odd :: Exp a -> Exp Bool
odd Exp a
n = Exp a
n Exp a -> Exp a -> Exp a
forall a. Integral a => a -> a -> a
`rem` Exp a
2 Exp a -> Exp a -> Exp Bool
forall a. Eq a => Exp a -> Exp a -> Exp Bool
/= Exp a
0

-- | @'gcd' x y@ is the non-negative factor of both @x@ and @y@ of which every
-- common factor of both @x@ and @y@ is also a factor; for example:
--
-- > gcd 4 2 = 2
-- > gcd (-4) 6 = 2
-- > gcd 0 4 = 4
-- > gcd 0 0 = 0
--
-- That is, the common divisor that is \"greatest\" in the divisibility
-- preordering.
--
gcd :: Integral a => Exp a -> Exp a -> Exp a
gcd :: Exp a -> Exp a -> Exp a
gcd Exp a
x Exp a
y = Exp a -> Exp a -> Exp a
forall a. Integral a => Exp a -> Exp a -> Exp a
gcd' (Exp a -> Exp a
forall a. Num a => a -> a
abs Exp a
x) (Exp a -> Exp a
forall a. Num a => a -> a
abs Exp a
y)
  where
    gcd' :: Integral a => Exp a -> Exp a -> Exp a
    gcd' :: Exp a -> Exp a -> Exp a
gcd' Exp a
u Exp a
v =
      let T2 Exp a
r Exp a
_ = (Exp (a, a) -> Exp Bool)
-> (Exp (a, a) -> Exp (a, a)) -> Exp (a, a) -> Exp (a, a)
forall e.
Elt e =>
(Exp e -> Exp Bool) -> (Exp e -> Exp e) -> Exp e -> Exp e
while (\(T2 Exp a
_ Exp a
b) -> Exp a
b Exp a -> Exp a -> Exp Bool
forall a. Eq a => Exp a -> Exp a -> Exp Bool
/= Exp a
0)
                         (\(T2 Exp a
a Exp a
b) -> Exp a -> Exp a -> Exp (a, a)
forall (con :: * -> *) x0 x1.
IsPattern con (x0, x1) (con x0, con x1) =>
con x0 -> con x1 -> con (x0, x1)
T2 Exp a
b (Exp a
a Exp a -> Exp a -> Exp a
forall a. Integral a => a -> a -> a
`rem` Exp a
b))
                         (Exp a -> Exp a -> Exp (a, a)
forall (con :: * -> *) x0 x1.
IsPattern con (x0, x1) (con x0, con x1) =>
con x0 -> con x1 -> con (x0, x1)
T2 Exp a
u Exp a
v)
      in Exp a
r


-- | @'lcm' x y@ is the smallest positive integer that both @x@ and @y@ divide.
--
lcm :: Integral a => Exp a -> Exp a -> Exp a
lcm :: Exp a -> Exp a -> Exp a
lcm Exp a
x Exp a
y
  = Exp Bool -> Exp a -> Exp a -> Exp a
forall t. Elt t => Exp Bool -> Exp t -> Exp t -> Exp t
cond (Exp a
x Exp a -> Exp a -> Exp Bool
forall a. Eq a => Exp a -> Exp a -> Exp Bool
== Exp a
0 Exp Bool -> Exp Bool -> Exp Bool
|| Exp a
y Exp a -> Exp a -> Exp Bool
forall a. Eq a => Exp a -> Exp a -> Exp Bool
== Exp a
0) Exp a
0
  (Exp a -> Exp a) -> Exp a -> Exp a
forall a b. (a -> b) -> a -> b
$ Exp a -> Exp a
forall a. Num a => a -> a
abs ((Exp a
x Exp a -> Exp a -> Exp a
forall a. Integral a => a -> a -> a
`quot` (Exp a -> Exp a -> Exp a
forall a. Integral a => Exp a -> Exp a -> Exp a
gcd Exp a
x Exp a
y)) Exp a -> Exp a -> Exp a
forall a. Num a => a -> a -> a
* Exp a
y)


-- | Raise a number to a non-negative integral power
--
infixr 8 ^
(^) :: forall a b. (Num a, Integral b) => Exp a -> Exp b -> Exp a
Exp a
x0 ^ :: Exp a -> Exp b -> Exp a
^ Exp b
y0 = Exp Bool -> Exp a -> Exp a -> Exp a
forall t. Elt t => Exp Bool -> Exp t -> Exp t -> Exp t
cond (Exp b
y0 Exp b -> Exp b -> Exp Bool
forall a. Ord a => Exp a -> Exp a -> Exp Bool
<= Exp b
0) Exp a
1 (Exp a -> Exp b -> Exp a
f Exp a
x0 Exp b
y0)
  where
    f :: Exp a -> Exp b -> Exp a
    f :: Exp a -> Exp b -> Exp a
f Exp a
x Exp b
y =
      let T2 Exp a
x' Exp b
y' = (Exp (a, b) -> Exp Bool)
-> (Exp (a, b) -> Exp (a, b)) -> Exp (a, b) -> Exp (a, b)
forall e.
Elt e =>
(Exp e -> Exp Bool) -> (Exp e -> Exp e) -> Exp e -> Exp e
while (\(T2 Exp a
_ Exp b
v) -> Exp b -> Exp Bool
forall a. Integral a => Exp a -> Exp Bool
even Exp b
v)
                           (\(T2 Exp a
u Exp b
v) -> Exp a -> Exp b -> Exp (a, b)
forall (con :: * -> *) x0 x1.
IsPattern con (x0, x1) (con x0, con x1) =>
con x0 -> con x1 -> con (x0, x1)
T2 (Exp a
u Exp a -> Exp a -> Exp a
forall a. Num a => a -> a -> a
* Exp a
u) (Exp b
v Exp b -> Exp b -> Exp b
forall a. Integral a => a -> a -> a
`quot` Exp b
2))
                           (Exp a -> Exp b -> Exp (a, b)
forall (con :: * -> *) x0 x1.
IsPattern con (x0, x1) (con x0, con x1) =>
con x0 -> con x1 -> con (x0, x1)
T2 Exp a
x Exp b
y)
      in
      Exp Bool -> Exp a -> Exp a -> Exp a
forall t. Elt t => Exp Bool -> Exp t -> Exp t -> Exp t
cond (Exp b
y' Exp b -> Exp b -> Exp Bool
forall a. Eq a => Exp a -> Exp a -> Exp Bool
== Exp b
1) Exp a
x' (Exp a -> Exp b -> Exp a -> Exp a
g (Exp a
x'Exp a -> Exp a -> Exp a
forall a. Num a => a -> a -> a
*Exp a
x') ((Exp b
y'Exp b -> Exp b -> Exp b
forall a. Num a => a -> a -> a
-Exp b
1) Exp b -> Exp b -> Exp b
forall a. Integral a => a -> a -> a
`quot` Exp b
2) Exp a
x')

    g :: Exp a -> Exp b -> Exp a -> Exp a
    g :: Exp a -> Exp b -> Exp a -> Exp a
g Exp a
x Exp b
y Exp a
z =
      let T3 Exp a
x' Exp b
_ Exp a
z' = (Exp (a, b, a) -> Exp Bool)
-> (Exp (a, b, a) -> Exp (a, b, a))
-> Exp (a, b, a)
-> Exp (a, b, a)
forall e.
Elt e =>
(Exp e -> Exp Bool) -> (Exp e -> Exp e) -> Exp e -> Exp e
while (\(T3 Exp a
_ Exp b
v Exp a
_) -> Exp b
v Exp b -> Exp b -> Exp Bool
forall a. Eq a => Exp a -> Exp a -> Exp Bool
/= Exp b
1)
                             (\(T3 Exp a
u Exp b
v Exp a
w) ->
                               Exp Bool -> Exp (a, b, a) -> Exp (a, b, a) -> Exp (a, b, a)
forall t. Elt t => Exp Bool -> Exp t -> Exp t -> Exp t
cond (Exp b -> Exp Bool
forall a. Integral a => Exp a -> Exp Bool
even Exp b
v) (Exp a -> Exp b -> Exp a -> Exp (a, b, a)
forall (con :: * -> *) x0 x1 x2.
IsPattern con (x0, x1, x2) (con x0, con x1, con x2) =>
con x0 -> con x1 -> con x2 -> con (x0, x1, x2)
T3 (Exp a
uExp a -> Exp a -> Exp a
forall a. Num a => a -> a -> a
*Exp a
u) (Exp b
v     Exp b -> Exp b -> Exp b
forall a. Integral a => a -> a -> a
`quot` Exp b
2) Exp a
w)
                                             (Exp a -> Exp b -> Exp a -> Exp (a, b, a)
forall (con :: * -> *) x0 x1 x2.
IsPattern con (x0, x1, x2) (con x0, con x1, con x2) =>
con x0 -> con x1 -> con x2 -> con (x0, x1, x2)
T3 (Exp a
uExp a -> Exp a -> Exp a
forall a. Num a => a -> a -> a
*Exp a
u) ((Exp b
vExp b -> Exp b -> Exp b
forall a. Num a => a -> a -> a
-Exp b
1) Exp b -> Exp b -> Exp b
forall a. Integral a => a -> a -> a
`quot` Exp b
2) (Exp a
wExp a -> Exp a -> Exp a
forall a. Num a => a -> a -> a
*Exp a
u)))
                             (Exp a -> Exp b -> Exp a -> Exp (a, b, a)
forall (con :: * -> *) x0 x1 x2.
IsPattern con (x0, x1, x2) (con x0, con x1, con x2) =>
con x0 -> con x1 -> con x2 -> con (x0, x1, x2)
T3 Exp a
x Exp b
y Exp a
z)
      in
      Exp a
x' Exp a -> Exp a -> Exp a
forall a. Num a => a -> a -> a
* Exp a
z'

-- | Raise a number to an integral power
--
infixr 8 ^^
(^^) :: (Fractional a, Integral b) => Exp a -> Exp b -> Exp a
Exp a
x ^^ :: Exp a -> Exp b -> Exp a
^^ Exp b
n
  = Exp Bool -> Exp a -> Exp a -> Exp a
forall t. Elt t => Exp Bool -> Exp t -> Exp t -> Exp t
cond (Exp b
n Exp b -> Exp b -> Exp Bool
forall a. Ord a => Exp a -> Exp a -> Exp Bool
>= Exp b
0)
  {- then -} (Exp a
x Exp a -> Exp b -> Exp a
forall a b. (Num a, Integral b) => Exp a -> Exp b -> Exp a
^ Exp b
n)
  {- else -} (Exp a -> Exp a
forall a. Fractional a => a -> a
recip (Exp a
x Exp a -> Exp b -> Exp a
forall a b. (Num a, Integral b) => Exp a -> Exp b -> Exp a
^ (Exp b -> Exp b
forall a. Num a => a -> a
negate Exp b
n)))


-- Conversions
-- -----------

-- |Convert a character to an 'Int'.
--
ord :: Exp Char -> Exp Int
ord :: Exp Char -> Exp Int
ord = Exp Char -> Exp Int
forall a b.
(Elt a, Elt b, IsIntegral (EltR a), IsNum (EltR b)) =>
Exp a -> Exp b
mkFromIntegral

-- |Convert an 'Int' into a character.
--
chr :: Exp Int -> Exp Char
chr :: Exp Int -> Exp Char
chr = Exp Int -> Exp Char
forall a b.
(Elt a, Elt b, IsIntegral (EltR a), IsNum (EltR b)) =>
Exp a -> Exp b
mkFromIntegral

-- |Convert a Boolean value to an 'Int', where 'False' turns into '0' and 'True'
-- into '1'.
--
boolToInt :: Exp Bool -> Exp Int
boolToInt :: Exp Bool -> Exp Int
boolToInt = Exp PrimBool -> Exp Int
forall a b.
(Elt a, Elt b, IsIntegral (EltR a), IsNum (EltR b)) =>
Exp a -> Exp b
mkFromIntegral (Exp PrimBool -> Exp Int)
-> (Exp Bool -> Exp PrimBool) -> Exp Bool -> Exp Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Coerce (EltR Bool) (EltR PrimBool) => Exp Bool -> Exp PrimBool
forall a b. Coerce (EltR a) (EltR b) => Exp a -> Exp b
mkCoerce @_ @Word8

-- |Reinterpret a value as another type. The two representations must have the
-- same bit size.
--
bitcast
    :: (Elt a, Elt b, IsScalar (EltR a), IsScalar (EltR b), BitSizeEq (EltR a) (EltR b))
    => Exp a
    -> Exp b
bitcast :: Exp a -> Exp b
bitcast = Exp a -> Exp b
forall b a.
(Elt a, Elt b, IsScalar (EltR a), IsScalar (EltR b),
 BitSizeEq (EltR a) (EltR b)) =>
Exp a -> Exp b
mkBitcast