{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE MagicHash           #-}
{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE PatternGuards       #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE RecordWildCards     #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications    #-}
{-# LANGUAGE TypeFamilies        #-}
{-# LANGUAGE TypeOperators       #-}
{-# LANGUAGE ViewPatterns        #-}
{-# OPTIONS_GHC -fno-warn-name-shadowing #-}
{-# OPTIONS_HADDOCK prune #-}
-- |
-- Module      : Data.Array.Accelerate.Interpreter
-- Description : Reference backend (interpreted)
-- Copyright   : [2008..2020] The Accelerate Team
-- License     : BSD3
--
-- Maintainer  : Trevor L. McDonell <trevor.mcdonell@gmail.com>
-- Stability   : experimental
-- Portability : non-portable (GHC extensions)
--
-- This interpreter is meant to be a reference implementation of the
-- semantics of the embedded array language. The emphasis is on defining
-- the semantics clearly, not on performance.
--

module Data.Array.Accelerate.Interpreter (

  Smart.Acc, Sugar.Arrays,
  Afunction, AfunctionR,

  -- * Interpret an array expression
  run, run1, runN,

  -- Internal (hidden)
  evalPrim, evalPrimConst, evalCoerceScalar, atraceOp,

) where

import Data.Array.Accelerate.AST                                    hiding ( Boundary(..) )
import Data.Array.Accelerate.AST.Environment
import Data.Array.Accelerate.AST.Var
import Data.Array.Accelerate.Array.Data
import Data.Array.Accelerate.Error
import Data.Array.Accelerate.Representation.Array
import Data.Array.Accelerate.Representation.Elt
import Data.Array.Accelerate.Representation.Shape
import Data.Array.Accelerate.Representation.Slice
import Data.Array.Accelerate.Representation.Stencil
import Data.Array.Accelerate.Representation.Tag
import Data.Array.Accelerate.Representation.Type
import Data.Array.Accelerate.Representation.Vec
import Data.Array.Accelerate.Trafo
import Data.Array.Accelerate.Trafo.Delayed                          ( DelayedOpenAfun, DelayedOpenAcc )
import Data.Array.Accelerate.Trafo.Sharing                          ( AfunctionR, AfunctionRepr(..), afunctionRepr )
import Data.Array.Accelerate.Type
import Data.Primitive.Vec
import qualified Data.Array.Accelerate.AST                          as AST
import qualified Data.Array.Accelerate.Debug.Internal.Flags         as Debug
import qualified Data.Array.Accelerate.Debug.Internal.Graph         as Debug
import qualified Data.Array.Accelerate.Debug.Internal.Stats         as Debug
import qualified Data.Array.Accelerate.Debug.Internal.Timed         as Debug
import qualified Data.Array.Accelerate.Smart                        as Smart
import qualified Data.Array.Accelerate.Sugar.Array                  as Sugar
import qualified Data.Array.Accelerate.Sugar.Elt                    as Sugar
import qualified Data.Array.Accelerate.Trafo.Delayed                as AST

import Control.DeepSeq
import Control.Exception
import Control.Monad
import Control.Monad.ST
import Data.Bits
import Data.Primitive.ByteArray
import Data.Primitive.Types
import Data.Text.Lazy.Builder
import Formatting
import System.IO
import System.IO.Unsafe                                             ( unsafePerformIO )
import Unsafe.Coerce
import qualified Data.Text.IO                                       as T
import Prelude                                                      hiding ( (!!), sum )


-- Program execution
-- -----------------

-- | Run a complete embedded array program using the reference interpreter.
--
run :: (HasCallStack, Sugar.Arrays a) => Smart.Acc a -> a
run :: forall a. (HasCallStack, Arrays a) => Acc a -> a
run Acc a
a = IO a -> a
forall a. IO a -> a
unsafePerformIO IO a
execute
  where
    !acc :: DelayedAcc (ArraysR a)
acc    = Acc a -> DelayedAcc (ArraysR a)
forall arrs. Acc arrs -> DelayedAcc (ArraysR arrs)
convertAcc Acc a
a
    execute :: IO a
execute = do
      DelayedAcc (ArraysR a) -> IO ()
forall (m :: * -> *) g. (MonadIO m, PrettyGraph g) => g -> m ()
Debug.dumpGraph (DelayedAcc (ArraysR a) -> IO ())
-> DelayedAcc (ArraysR a) -> IO ()
forall a b. NFData a => (a -> b) -> a -> b
$!! DelayedAcc (ArraysR a)
acc
      IO ()
Debug.dumpSimplStats
      WithReprs (ArraysR a)
res <- Builder
-> Format Builder (Double -> Double -> Builder)
-> IO (WithReprs (ArraysR a))
-> IO (WithReprs (ArraysR a))
forall a.
Builder
-> Format Builder (Double -> Double -> Builder) -> IO a -> IO a
phase Builder
"execute" Format Builder (Double -> Double -> Builder)
forall r. Format r (Double -> Double -> r)
Debug.elapsed (IO (WithReprs (ArraysR a)) -> IO (WithReprs (ArraysR a)))
-> IO (WithReprs (ArraysR a)) -> IO (WithReprs (ArraysR a))
forall a b. (a -> b) -> a -> b
$ WithReprs (ArraysR a) -> IO (WithReprs (ArraysR a))
forall a. a -> IO a
evaluate (WithReprs (ArraysR a) -> IO (WithReprs (ArraysR a)))
-> WithReprs (ArraysR a) -> IO (WithReprs (ArraysR a))
forall a b. (a -> b) -> a -> b
$ DelayedAcc (ArraysR a) -> Val () -> WithReprs (ArraysR a)
forall aenv a.
HasCallStack =>
DelayedOpenAcc aenv a -> Val aenv -> WithReprs a
evalOpenAcc DelayedAcc (ArraysR a)
acc Val ()
Empty
      a -> IO a
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (a -> IO a) -> a -> IO a
forall a b. (a -> b) -> a -> b
$ ArraysR a -> a
forall a. Arrays a => ArraysR a -> a
Sugar.toArr (ArraysR a -> a) -> ArraysR a -> a
forall a b. (a -> b) -> a -> b
$ WithReprs (ArraysR a) -> ArraysR a
forall a b. (a, b) -> b
snd WithReprs (ArraysR a)
res

-- | This is 'runN' specialised to an array program of one argument.
--
run1 :: (HasCallStack, Sugar.Arrays a, Sugar.Arrays b) => (Smart.Acc a -> Smart.Acc b) -> a -> b
run1 :: forall a b.
(HasCallStack, Arrays a, Arrays b) =>
(Acc a -> Acc b) -> a -> b
run1 = (Acc a -> Acc b) -> AfunctionR (Acc a -> Acc b)
(Acc a -> Acc b) -> a -> b
forall f. (HasCallStack, Afunction f) => f -> AfunctionR f
runN

-- | Prepare and execute an embedded array program.
--
runN :: forall f. (HasCallStack, Afunction f) => f -> AfunctionR f
runN :: forall f. (HasCallStack, Afunction f) => f -> AfunctionR f
runN f
f = AfunctionR f
go
  where
    !acc :: DelayedAfun (ArraysFunctionR f)
acc    = f -> DelayedAfun (ArraysFunctionR f)
forall f. Afunction f => f -> DelayedAfun (ArraysFunctionR f)
convertAfun f
f
    !afun :: DelayedAfun (ArraysFunctionR f)
afun   = IO (DelayedAfun (ArraysFunctionR f))
-> DelayedAfun (ArraysFunctionR f)
forall a. IO a -> a
unsafePerformIO (IO (DelayedAfun (ArraysFunctionR f))
 -> DelayedAfun (ArraysFunctionR f))
-> IO (DelayedAfun (ArraysFunctionR f))
-> DelayedAfun (ArraysFunctionR f)
forall a b. (a -> b) -> a -> b
$ do
                DelayedAfun (ArraysFunctionR f) -> IO ()
forall (m :: * -> *) g. (MonadIO m, PrettyGraph g) => g -> m ()
Debug.dumpGraph (DelayedAfun (ArraysFunctionR f) -> IO ())
-> DelayedAfun (ArraysFunctionR f) -> IO ()
forall a b. NFData a => (a -> b) -> a -> b
$!! DelayedAfun (ArraysFunctionR f)
acc
                IO ()
Debug.dumpSimplStats
                DelayedAfun (ArraysFunctionR f)
-> IO (DelayedAfun (ArraysFunctionR f))
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return DelayedAfun (ArraysFunctionR f)
acc
    !go :: AfunctionR f
go     = AfunctionRepr f (AfunctionR f) (ArraysFunctionR f)
-> DelayedAfun (ArraysFunctionR f) -> Val () -> AfunctionR f
forall g aenv.
AfunctionRepr g (AfunctionR g) (ArraysFunctionR g)
-> DelayedOpenAfun aenv (ArraysFunctionR g)
-> Val aenv
-> AfunctionR g
eval (forall f.
(Afunction f, HasCallStack) =>
AfunctionRepr f (AfunctionR f) (ArraysFunctionR f)
afunctionRepr @f) DelayedAfun (ArraysFunctionR f)
afun Val ()
Empty
    --
    eval :: AfunctionRepr g (AfunctionR g) (ArraysFunctionR g)
         -> DelayedOpenAfun aenv (ArraysFunctionR g)
         -> Val aenv
         -> AfunctionR g
    eval :: forall g aenv.
AfunctionRepr g (AfunctionR g) (ArraysFunctionR g)
-> DelayedOpenAfun aenv (ArraysFunctionR g)
-> Val aenv
-> AfunctionR g
eval (AfunctionReprLam AfunctionRepr b br breprr
reprF) (Alam ALeftHandSide a aenv aenv'
lhs PreOpenAfun DelayedOpenAcc aenv' t1
f) Val aenv
aenv = \a
a -> AfunctionRepr b (AfunctionR b) (ArraysFunctionR b)
-> DelayedOpenAfun aenv' (ArraysFunctionR b)
-> Val aenv'
-> AfunctionR b
forall g aenv.
AfunctionRepr g (AfunctionR g) (ArraysFunctionR g)
-> DelayedOpenAfun aenv (ArraysFunctionR g)
-> Val aenv
-> AfunctionR g
eval AfunctionRepr b br breprr
AfunctionRepr b (AfunctionR b) (ArraysFunctionR b)
reprF PreOpenAfun DelayedOpenAcc aenv' t1
DelayedOpenAfun aenv' (ArraysFunctionR b)
f (Val aenv' -> AfunctionR b) -> Val aenv' -> AfunctionR b
forall a b. (a -> b) -> a -> b
$ Val aenv
aenv Val aenv -> (ALeftHandSide a aenv aenv', a) -> Val aenv'
forall env (s :: * -> *) t env'.
Val env -> (LeftHandSide s t env env', t) -> Val env'
`push` (ALeftHandSide a aenv aenv'
lhs, a -> ArraysR a
forall a. Arrays a => a -> ArraysR a
Sugar.fromArr a
a)
    eval AfunctionRepr g (AfunctionR g) (ArraysFunctionR g)
AfunctionReprBody        (Abody DelayedOpenAcc aenv (ArraysFunctionR g)
b)    Val aenv
aenv = IO (AfunctionR g) -> AfunctionR g
forall a. IO a -> a
unsafePerformIO (IO (AfunctionR g) -> AfunctionR g)
-> IO (AfunctionR g) -> AfunctionR g
forall a b. (a -> b) -> a -> b
$ Builder
-> Format Builder (Double -> Double -> Builder)
-> IO (AfunctionR g)
-> IO (AfunctionR g)
forall a.
Builder
-> Format Builder (Double -> Double -> Builder) -> IO a -> IO a
phase Builder
"execute" Format Builder (Double -> Double -> Builder)
forall r. Format r (Double -> Double -> r)
Debug.elapsed (ArraysR (AfunctionR g) -> AfunctionR g
forall a. Arrays a => ArraysR a -> a
Sugar.toArr (ArraysR (AfunctionR g) -> AfunctionR g)
-> ((ArraysR (ArraysR (AfunctionR g)), ArraysR (AfunctionR g))
    -> ArraysR (AfunctionR g))
-> (ArraysR (ArraysR (AfunctionR g)), ArraysR (AfunctionR g))
-> AfunctionR g
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ArraysR (ArraysR (AfunctionR g)), ArraysR (AfunctionR g))
-> ArraysR (AfunctionR g)
forall a b. (a, b) -> b
snd ((ArraysR (ArraysR (AfunctionR g)), ArraysR (AfunctionR g))
 -> AfunctionR g)
-> IO (ArraysR (ArraysR (AfunctionR g)), ArraysR (AfunctionR g))
-> IO (AfunctionR g)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (ArraysR (ArraysR (AfunctionR g)), ArraysR (AfunctionR g))
-> IO (ArraysR (ArraysR (AfunctionR g)), ArraysR (AfunctionR g))
forall a. a -> IO a
evaluate (DelayedOpenAcc aenv (ArraysR (AfunctionR g))
-> Val aenv
-> (ArraysR (ArraysR (AfunctionR g)), ArraysR (AfunctionR g))
forall aenv a.
HasCallStack =>
DelayedOpenAcc aenv a -> Val aenv -> WithReprs a
evalOpenAcc DelayedOpenAcc aenv (ArraysR (AfunctionR g))
DelayedOpenAcc aenv (ArraysFunctionR g)
b Val aenv
aenv))
    eval AfunctionRepr g (AfunctionR g) (ArraysFunctionR g)
_                        PreOpenAfun DelayedOpenAcc aenv (ArraysFunctionR g)
_aenv        Val aenv
_    = [Char] -> AfunctionR g
forall a. HasCallStack => [Char] -> a
error [Char]
"Two men say they're Jesus; one of them must be wrong"

-- -- | Stream a lazily read list of input arrays through the given program,
-- -- collecting results as we go
-- --
-- streamOut :: Arrays a => Sugar.Seq [a] -> [a]
-- streamOut seq = let seq' = convertSeqWith config seq
--                 in evalDelayedSeq defaultSeqConfig seq'


-- Debugging
-- ---------

phase :: Builder -> Format Builder (Double -> Double -> Builder) -> IO a -> IO a
phase :: forall a.
Builder
-> Format Builder (Double -> Double -> Builder) -> IO a -> IO a
phase Builder
n Format Builder (Double -> Double -> Builder)
fmt IO a
go = Flag
-> Format Builder (Double -> Double -> Builder) -> IO a -> IO a
forall (m :: * -> *) a.
MonadIO m =>
Flag -> Format Builder (Double -> Double -> Builder) -> m a -> m a
Debug.timed Flag
Debug.dump_phases (Builder
-> Format
     (Double -> Double -> Builder) (Double -> Double -> Builder)
forall r. Builder -> Format r r
now (Builder
"phase " Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Builder
n Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Builder
": ") Format (Double -> Double -> Builder) (Double -> Double -> Builder)
-> Format Builder (Double -> Double -> Builder)
-> Format Builder (Double -> Double -> Builder)
forall r a r'. Format r a -> Format r' r -> Format r' a
% Format Builder (Double -> Double -> Builder)
fmt) IO a
go


-- Delayed Arrays
-- --------------

-- Note that in contrast to the representation used in the optimised AST, the
-- delayed array representation used here is _only_ for delayed arrays --- we do
-- not require an optional Manifest|Delayed data type to evaluate the program.
--
data Delayed a where
  Delayed :: ArrayR (Array sh e)
          -> sh
          -> (sh -> e)
          -> (Int -> e)
          -> Delayed (Array sh e)


-- Array expression evaluation
-- ---------------------------

type WithReprs acc = (ArraysR acc, acc)

fromFunction' :: ArrayR (Array sh e) -> sh -> (sh -> e) -> WithReprs (Array sh e)
fromFunction' :: forall sh e.
ArrayR (Array sh e) -> sh -> (sh -> e) -> WithReprs (Array sh e)
fromFunction' ArrayR (Array sh e)
repr sh
sh sh -> e
f = (ArrayR (Array sh e) -> TupR ArrayR (Array sh e)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle ArrayR (Array sh e)
repr, ArrayR (Array sh e) -> sh -> (sh -> e) -> Array sh e
forall sh e. ArrayR (Array sh e) -> sh -> (sh -> e) -> Array sh e
fromFunction ArrayR (Array sh e)
repr sh
sh sh -> e
f)

-- Evaluate an open array function
--
evalOpenAfun :: HasCallStack => DelayedOpenAfun aenv f -> Val aenv -> f
evalOpenAfun :: forall aenv f.
HasCallStack =>
DelayedOpenAfun aenv f -> Val aenv -> f
evalOpenAfun (Alam ALeftHandSide a aenv aenv'
lhs PreOpenAfun DelayedOpenAcc aenv' t1
f) Val aenv
aenv = \a
a -> PreOpenAfun DelayedOpenAcc aenv' t1 -> Val aenv' -> t1
forall aenv f.
HasCallStack =>
DelayedOpenAfun aenv f -> Val aenv -> f
evalOpenAfun PreOpenAfun DelayedOpenAcc aenv' t1
f (Val aenv' -> t1) -> Val aenv' -> t1
forall a b. (a -> b) -> a -> b
$ Val aenv
aenv Val aenv -> (ALeftHandSide a aenv aenv', a) -> Val aenv'
forall env (s :: * -> *) t env'.
Val env -> (LeftHandSide s t env env', t) -> Val env'
`push` (ALeftHandSide a aenv aenv'
lhs, a
a)
evalOpenAfun (Abody DelayedOpenAcc aenv f
b)    Val aenv
aenv = (ArraysR f, f) -> f
forall a b. (a, b) -> b
snd ((ArraysR f, f) -> f) -> (ArraysR f, f) -> f
forall a b. (a -> b) -> a -> b
$ DelayedOpenAcc aenv f -> Val aenv -> (ArraysR f, f)
forall aenv a.
HasCallStack =>
DelayedOpenAcc aenv a -> Val aenv -> WithReprs a
evalOpenAcc DelayedOpenAcc aenv f
b Val aenv
aenv


-- The core interpreter for optimised array programs
--
evalOpenAcc
    :: forall aenv a. HasCallStack
    => DelayedOpenAcc aenv a
    -> Val aenv
    -> WithReprs a
evalOpenAcc :: forall aenv a.
HasCallStack =>
DelayedOpenAcc aenv a -> Val aenv -> WithReprs a
evalOpenAcc AST.Delayed{}       Val aenv
_    = Format (ArraysR (Array sh e), Array sh e) (WithReprs a)
-> WithReprs a
forall r a. HasCallStack => Format r a -> a
internalError Format (ArraysR (Array sh e), Array sh e) (WithReprs a)
"expected manifest array"
evalOpenAcc (AST.Manifest PreOpenAcc DelayedOpenAcc aenv a
pacc) Val aenv
aenv =
  let
      manifest :: forall a'. HasCallStack => DelayedOpenAcc aenv a' -> WithReprs a'
      manifest :: forall a'. HasCallStack => DelayedOpenAcc aenv a' -> WithReprs a'
manifest DelayedOpenAcc aenv a'
acc =
        let (ArraysR a'
repr, a'
a') = DelayedOpenAcc aenv a' -> Val aenv -> (ArraysR a', a')
forall aenv a.
HasCallStack =>
DelayedOpenAcc aenv a -> Val aenv -> WithReprs a
evalOpenAcc DelayedOpenAcc aenv a'
acc Val aenv
aenv
        in  ArraysR a' -> a' -> ()
forall arrs. ArraysR arrs -> arrs -> ()
rnfArraysR ArraysR a'
repr a'
a' () -> (ArraysR a', a') -> (ArraysR a', a')
forall a b. a -> b -> b
`seq` (ArraysR a'
repr, a'
a')

      delayed :: DelayedOpenAcc aenv (Array sh e) -> Delayed (Array sh e)
      delayed :: forall sh e.
DelayedOpenAcc aenv (Array sh e) -> Delayed (Array sh e)
delayed AST.Delayed{ArrayR (Array sh e)
Exp aenv sh
Fun aenv (sh -> e)
Fun aenv (Int -> e)
reprD :: ArrayR (Array sh e)
extentD :: Exp aenv sh
indexD :: Fun aenv (sh -> e)
linearIndexD :: Fun aenv (Int -> e)
reprD :: forall sh e aenv.
DelayedOpenAcc aenv (Array sh e) -> ArrayR (Array sh e)
extentD :: forall sh e aenv. DelayedOpenAcc aenv (Array sh e) -> Exp aenv sh
indexD :: forall sh e aenv.
DelayedOpenAcc aenv (Array sh e) -> Fun aenv (sh -> e)
linearIndexD :: forall sh e aenv.
DelayedOpenAcc aenv (Array sh e) -> Fun aenv (Int -> e)
..} = ArrayR (Array sh e)
-> sh -> (sh -> e) -> (Int -> e) -> Delayed (Array sh e)
forall sh e.
ArrayR (Array sh e)
-> sh -> (sh -> e) -> (Int -> e) -> Delayed (Array sh e)
Delayed ArrayR (Array sh e)
ArrayR (Array sh e)
reprD (Exp aenv sh -> sh
forall t. Exp aenv t -> t
evalE Exp aenv sh
Exp aenv sh
extentD) (Fun aenv (sh -> e) -> sh -> e
forall f. Fun aenv f -> f
evalF Fun aenv (sh -> e)
Fun aenv (sh -> e)
indexD) (Fun aenv (Int -> e) -> Int -> e
forall f. Fun aenv f -> f
evalF Fun aenv (Int -> e)
Fun aenv (Int -> e)
linearIndexD)
      delayed DelayedOpenAcc aenv (Array sh e)
a' = ArrayR (Array sh e)
-> sh -> (sh -> e) -> (Int -> e) -> Delayed (Array sh e)
forall sh e.
ArrayR (Array sh e)
-> sh -> (sh -> e) -> (Int -> e) -> Delayed (Array sh e)
Delayed ArrayR (Array sh e)
aR (Array sh e -> sh
forall sh e. Array sh e -> sh
shape Array sh e
a) (ArrayR (Array sh e) -> Array sh e -> sh -> e
forall sh e. ArrayR (Array sh e) -> Array sh e -> sh -> e
indexArray ArrayR (Array sh e)
aR Array sh e
a) (TypeR e -> Array sh e -> Int -> e
forall e sh. TypeR e -> Array sh e -> Int -> e
linearIndexArray (ArrayR (Array sh e) -> TypeR e
forall sh e. ArrayR (Array sh e) -> TypeR e
arrayRtype ArrayR (Array sh e)
aR) Array sh e
a)
        where
          (TupRsingle ArrayR (Array sh e)
aR, Array sh e
a) = DelayedOpenAcc aenv (Array sh e)
-> (ArraysR (Array sh e), Array sh e)
forall a'. HasCallStack => DelayedOpenAcc aenv a' -> WithReprs a'
manifest DelayedOpenAcc aenv (Array sh e)
a'

      evalE :: Exp aenv t -> t
      evalE :: forall t. Exp aenv t -> t
evalE Exp aenv t
exp = Exp aenv t -> Val aenv -> t
forall aenv t. HasCallStack => Exp aenv t -> Val aenv -> t
evalExp Exp aenv t
exp Val aenv
aenv

      evalF :: Fun aenv f -> f
      evalF :: forall f. Fun aenv f -> f
evalF Fun aenv f
fun = Fun aenv f -> Val aenv -> f
forall aenv t. HasCallStack => Fun aenv t -> Val aenv -> t
evalFun Fun aenv f
fun Val aenv
aenv

      evalB :: AST.Boundary aenv t -> Boundary t
      evalB :: forall t. Boundary aenv t -> Boundary t
evalB Boundary aenv t
bnd = Boundary aenv t -> Val aenv -> Boundary t
forall aenv t.
HasCallStack =>
Boundary aenv t -> Val aenv -> Boundary t
evalBoundary Boundary aenv t
bnd Val aenv
aenv

      dir :: Direction -> t -> t -> t
      dir :: forall t. Direction -> t -> t -> t
dir Direction
LeftToRight t
l t
_ = t
l
      dir Direction
RightToLeft t
_ t
r = t
r
  in
  case PreOpenAcc DelayedOpenAcc aenv a
pacc of
    Avar (Var ArrayR (Array sh e)
repr Idx aenv (Array sh e)
ix)            -> (ArrayR a -> TupR ArrayR a
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle ArrayR a
ArrayR (Array sh e)
repr, Idx aenv a -> Val aenv -> a
forall env t. Idx env t -> Val env -> t
prj Idx aenv a
Idx aenv (Array sh e)
ix Val aenv
aenv)
    Alet ALeftHandSide bndArrs aenv aenv'
lhs DelayedOpenAcc aenv bndArrs
acc1 DelayedOpenAcc aenv' a
acc2            -> DelayedOpenAcc aenv' a -> Val aenv' -> WithReprs a
forall aenv a.
HasCallStack =>
DelayedOpenAcc aenv a -> Val aenv -> WithReprs a
evalOpenAcc DelayedOpenAcc aenv' a
acc2 (Val aenv' -> WithReprs a) -> Val aenv' -> WithReprs a
forall a b. (a -> b) -> a -> b
$ Val aenv
aenv Val aenv
-> (ALeftHandSide bndArrs aenv aenv', bndArrs) -> Val aenv'
forall env (s :: * -> *) t env'.
Val env -> (LeftHandSide s t env env', t) -> Val env'
`push` (ALeftHandSide bndArrs aenv aenv'
lhs, (ArraysR bndArrs, bndArrs) -> bndArrs
forall a b. (a, b) -> b
snd ((ArraysR bndArrs, bndArrs) -> bndArrs)
-> (ArraysR bndArrs, bndArrs) -> bndArrs
forall a b. (a -> b) -> a -> b
$ DelayedOpenAcc aenv bndArrs -> (ArraysR bndArrs, bndArrs)
forall a'. HasCallStack => DelayedOpenAcc aenv a' -> WithReprs a'
manifest DelayedOpenAcc aenv bndArrs
acc1)
    Apair DelayedOpenAcc aenv as
acc1 DelayedOpenAcc aenv bs
acc2               -> let (ArraysR as
r1, as
a1) = DelayedOpenAcc aenv as -> (ArraysR as, as)
forall a'. HasCallStack => DelayedOpenAcc aenv a' -> WithReprs a'
manifest DelayedOpenAcc aenv as
acc1
                                         (ArraysR bs
r2, bs
a2) = DelayedOpenAcc aenv bs -> (ArraysR bs, bs)
forall a'. HasCallStack => DelayedOpenAcc aenv a' -> WithReprs a'
manifest DelayedOpenAcc aenv bs
acc2
                                     in
                                     (ArraysR as -> ArraysR bs -> TupR ArrayR (as, bs)
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
TupRpair ArraysR as
r1 ArraysR bs
r2, (as
a1, bs
a2))
    PreOpenAcc DelayedOpenAcc aenv a
Anil                          -> (TupR ArrayR a
TupR ArrayR ()
forall (s :: * -> *). TupR s ()
TupRunit, ())
    Atrace Message arrs1
msg DelayedOpenAcc aenv arrs1
as DelayedOpenAcc aenv a
bs              -> IO (WithReprs a) -> WithReprs a
forall a. IO a -> a
unsafePerformIO (IO (WithReprs a) -> WithReprs a)
-> IO (WithReprs a) -> WithReprs a
forall a b. (a -> b) -> a -> b
$ DelayedOpenAcc aenv a -> WithReprs a
forall a'. HasCallStack => DelayedOpenAcc aenv a' -> WithReprs a'
manifest DelayedOpenAcc aenv a
bs WithReprs a -> IO () -> IO (WithReprs a)
forall a b. a -> IO b -> IO a
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Message arrs1 -> arrs1 -> IO ()
forall as. Message as -> as -> IO ()
atraceOp Message arrs1
msg ((ArraysR arrs1, arrs1) -> arrs1
forall a b. (a, b) -> b
snd ((ArraysR arrs1, arrs1) -> arrs1)
-> (ArraysR arrs1, arrs1) -> arrs1
forall a b. (a -> b) -> a -> b
$ DelayedOpenAcc aenv arrs1 -> (ArraysR arrs1, arrs1)
forall a'. HasCallStack => DelayedOpenAcc aenv a' -> WithReprs a'
manifest DelayedOpenAcc aenv arrs1
as)
    Apply TupR ArrayR a
repr PreOpenAfun DelayedOpenAcc aenv (arrs1 -> a)
afun DelayedOpenAcc aenv arrs1
acc           -> (TupR ArrayR a
repr, PreOpenAfun DelayedOpenAcc aenv (arrs1 -> a)
-> Val aenv -> arrs1 -> a
forall aenv f.
HasCallStack =>
DelayedOpenAfun aenv f -> Val aenv -> f
evalOpenAfun PreOpenAfun DelayedOpenAcc aenv (arrs1 -> a)
afun Val aenv
aenv (arrs1 -> a) -> arrs1 -> a
forall a b. (a -> b) -> a -> b
$ (ArraysR arrs1, arrs1) -> arrs1
forall a b. (a, b) -> b
snd ((ArraysR arrs1, arrs1) -> arrs1)
-> (ArraysR arrs1, arrs1) -> arrs1
forall a b. (a -> b) -> a -> b
$ DelayedOpenAcc aenv arrs1 -> (ArraysR arrs1, arrs1)
forall a'. HasCallStack => DelayedOpenAcc aenv a' -> WithReprs a'
manifest DelayedOpenAcc aenv arrs1
acc)
    Aforeign TupR ArrayR a
repr asm (as -> a)
_ PreAfun DelayedOpenAcc (as -> a)
afun DelayedOpenAcc aenv as
acc      -> (TupR ArrayR a
repr, PreAfun DelayedOpenAcc (as -> a) -> Val () -> as -> a
forall aenv f.
HasCallStack =>
DelayedOpenAfun aenv f -> Val aenv -> f
evalOpenAfun PreAfun DelayedOpenAcc (as -> a)
afun Val ()
Empty (as -> a) -> as -> a
forall a b. (a -> b) -> a -> b
$ (ArraysR as, as) -> as
forall a b. (a, b) -> b
snd ((ArraysR as, as) -> as) -> (ArraysR as, as) -> as
forall a b. (a -> b) -> a -> b
$ DelayedOpenAcc aenv as -> (ArraysR as, as)
forall a'. HasCallStack => DelayedOpenAcc aenv a' -> WithReprs a'
manifest DelayedOpenAcc aenv as
acc)
    Acond Exp aenv PrimBool
p DelayedOpenAcc aenv a
acc1 DelayedOpenAcc aenv a
acc2
      | PrimBool -> Bool
toBool (Exp aenv PrimBool -> PrimBool
forall t. Exp aenv t -> t
evalE Exp aenv PrimBool
p)          -> DelayedOpenAcc aenv a -> WithReprs a
forall a'. HasCallStack => DelayedOpenAcc aenv a' -> WithReprs a'
manifest DelayedOpenAcc aenv a
acc1
      | Bool
otherwise                 -> DelayedOpenAcc aenv a -> WithReprs a
forall a'. HasCallStack => DelayedOpenAcc aenv a' -> WithReprs a'
manifest DelayedOpenAcc aenv a
acc2

    Awhile PreOpenAfun DelayedOpenAcc aenv (a -> Scalar PrimBool)
cond PreOpenAfun DelayedOpenAcc aenv (a -> a)
body DelayedOpenAcc aenv a
acc          -> (TupR ArrayR a
repr, a -> a
go a
initial)
      where
        (TupR ArrayR a
repr, a
initial) = DelayedOpenAcc aenv a -> WithReprs a
forall a'. HasCallStack => DelayedOpenAcc aenv a' -> WithReprs a'
manifest DelayedOpenAcc aenv a
acc
        p :: a -> Scalar PrimBool
p               = PreOpenAfun DelayedOpenAcc aenv (a -> Scalar PrimBool)
-> Val aenv -> a -> Scalar PrimBool
forall aenv f.
HasCallStack =>
DelayedOpenAfun aenv f -> Val aenv -> f
evalOpenAfun PreOpenAfun DelayedOpenAcc aenv (a -> Scalar PrimBool)
cond Val aenv
aenv
        f :: a -> a
f               = PreOpenAfun DelayedOpenAcc aenv (a -> a) -> Val aenv -> a -> a
forall aenv f.
HasCallStack =>
DelayedOpenAfun aenv f -> Val aenv -> f
evalOpenAfun PreOpenAfun DelayedOpenAcc aenv (a -> a)
body Val aenv
aenv
        go :: a -> a
go !a
x
          | PrimBool -> Bool
toBool (TypeR PrimBool -> Scalar PrimBool -> Int -> PrimBool
forall e sh. TypeR e -> Array sh e -> Int -> e
linearIndexArray (forall a. Elt a => TypeR (EltR a)
Sugar.eltR @Word8) (a -> Scalar PrimBool
p a
x) Int
0) = a -> a
go (a -> a
f a
x)
          | Bool
otherwise                                             = a
x

    Use ArrayR (Array sh e)
repr Array sh e
arr                  -> (ArrayR a -> TupR ArrayR a
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle ArrayR a
ArrayR (Array sh e)
repr, a
Array sh e
arr)
    Unit TypeR e
tp Exp aenv e
e                     -> TypeR e -> e -> WithReprs (Array () e)
forall e. TypeR e -> e -> WithReprs (Scalar e)
unitOp TypeR e
tp (Exp aenv e -> e
forall t. Exp aenv t -> t
evalE Exp aenv e
e)
    -- Collect s                     -> evalSeq defaultSeqConfig s aenv

    -- Producers
    -- ---------
    Map TypeR e'
tp Fun aenv (e -> e')
f DelayedOpenAcc aenv (Array sh e)
acc                  -> TypeR e'
-> (e -> e') -> Delayed (Array sh e) -> WithReprs (Array sh e')
forall b a sh.
TypeR b
-> (a -> b) -> Delayed (Array sh a) -> WithReprs (Array sh b)
mapOp TypeR e'
tp (Fun aenv (e -> e') -> e -> e'
forall f. Fun aenv f -> f
evalF Fun aenv (e -> e')
f) (DelayedOpenAcc aenv (Array sh e) -> Delayed (Array sh e)
forall sh e.
DelayedOpenAcc aenv (Array sh e) -> Delayed (Array sh e)
delayed DelayedOpenAcc aenv (Array sh e)
acc)
    Generate ArrayR (Array sh e)
repr Exp aenv sh
sh Fun aenv (sh -> e)
f            -> ArrayR (Array sh e) -> sh -> (sh -> e) -> WithReprs (Array sh e)
forall sh e.
ArrayR (Array sh e) -> sh -> (sh -> e) -> WithReprs (Array sh e)
generateOp ArrayR (Array sh e)
repr (Exp aenv sh -> sh
forall t. Exp aenv t -> t
evalE Exp aenv sh
sh) (Fun aenv (sh -> e) -> sh -> e
forall f. Fun aenv f -> f
evalF Fun aenv (sh -> e)
f)
    Transform ArrayR (Array sh' b)
repr Exp aenv sh'
sh Fun aenv (sh' -> sh)
p Fun aenv (a1 -> b)
f DelayedOpenAcc aenv (Array sh a1)
acc     -> ArrayR (Array sh' b)
-> sh'
-> (sh' -> sh)
-> (a1 -> b)
-> Delayed (Array sh a1)
-> WithReprs (Array sh' b)
forall sh' b sh a.
ArrayR (Array sh' b)
-> sh'
-> (sh' -> sh)
-> (a -> b)
-> Delayed (Array sh a)
-> WithReprs (Array sh' b)
transformOp ArrayR (Array sh' b)
repr (Exp aenv sh' -> sh'
forall t. Exp aenv t -> t
evalE Exp aenv sh'
sh) (Fun aenv (sh' -> sh) -> sh' -> sh
forall f. Fun aenv f -> f
evalF Fun aenv (sh' -> sh)
p) (Fun aenv (a1 -> b) -> a1 -> b
forall f. Fun aenv f -> f
evalF Fun aenv (a1 -> b)
f) (DelayedOpenAcc aenv (Array sh a1) -> Delayed (Array sh a1)
forall sh e.
DelayedOpenAcc aenv (Array sh e) -> Delayed (Array sh e)
delayed DelayedOpenAcc aenv (Array sh a1)
acc)
    Backpermute ShapeR sh'
shr Exp aenv sh'
sh Fun aenv (sh' -> sh)
p DelayedOpenAcc aenv (Array sh e)
acc      -> ShapeR sh'
-> sh'
-> (sh' -> sh)
-> Delayed (Array sh e)
-> WithReprs (Array sh' e)
forall sh' sh e.
ShapeR sh'
-> sh'
-> (sh' -> sh)
-> Delayed (Array sh e)
-> WithReprs (Array sh' e)
backpermuteOp ShapeR sh'
shr (Exp aenv sh' -> sh'
forall t. Exp aenv t -> t
evalE Exp aenv sh'
sh) (Fun aenv (sh' -> sh) -> sh' -> sh
forall f. Fun aenv f -> f
evalF Fun aenv (sh' -> sh)
p) (DelayedOpenAcc aenv (Array sh e) -> Delayed (Array sh e)
forall sh e.
DelayedOpenAcc aenv (Array sh e) -> Delayed (Array sh e)
delayed DelayedOpenAcc aenv (Array sh e)
acc)
    Reshape ShapeR sh
shr Exp aenv sh
sh DelayedOpenAcc aenv (Array sh' e)
acc            -> ShapeR sh
-> sh -> WithReprs (Array sh' e) -> WithReprs (Array sh e)
forall sh sh' e.
HasCallStack =>
ShapeR sh
-> sh -> WithReprs (Array sh' e) -> WithReprs (Array sh e)
reshapeOp ShapeR sh
shr (Exp aenv sh -> sh
forall t. Exp aenv t -> t
evalE Exp aenv sh
sh) (DelayedOpenAcc aenv (Array sh' e) -> WithReprs (Array sh' e)
forall a'. HasCallStack => DelayedOpenAcc aenv a' -> WithReprs a'
manifest DelayedOpenAcc aenv (Array sh' e)
acc)

    ZipWith TypeR e3
tp Fun aenv (e1 -> e2 -> e3)
f DelayedOpenAcc aenv (Array sh e1)
acc1 DelayedOpenAcc aenv (Array sh e2)
acc2        -> TypeR e3
-> (e1 -> e2 -> e3)
-> Delayed (Array sh e1)
-> Delayed (Array sh e2)
-> WithReprs (Array sh e3)
forall c a b sh.
TypeR c
-> (a -> b -> c)
-> Delayed (Array sh a)
-> Delayed (Array sh b)
-> WithReprs (Array sh c)
zipWithOp TypeR e3
tp (Fun aenv (e1 -> e2 -> e3) -> e1 -> e2 -> e3
forall f. Fun aenv f -> f
evalF Fun aenv (e1 -> e2 -> e3)
f) (DelayedOpenAcc aenv (Array sh e1) -> Delayed (Array sh e1)
forall sh e.
DelayedOpenAcc aenv (Array sh e) -> Delayed (Array sh e)
delayed DelayedOpenAcc aenv (Array sh e1)
acc1) (DelayedOpenAcc aenv (Array sh e2) -> Delayed (Array sh e2)
forall sh e.
DelayedOpenAcc aenv (Array sh e) -> Delayed (Array sh e)
delayed DelayedOpenAcc aenv (Array sh e2)
acc2)
    Replicate SliceIndex slix sl co sh
slice Exp aenv slix
slix DelayedOpenAcc aenv (Array sl e)
acc      -> SliceIndex slix sl co sh
-> slix -> WithReprs (Array sl e) -> WithReprs (Array sh e)
forall slix sl co sh e.
SliceIndex slix sl co sh
-> slix -> WithReprs (Array sl e) -> WithReprs (Array sh e)
replicateOp SliceIndex slix sl co sh
slice (Exp aenv slix -> slix
forall t. Exp aenv t -> t
evalE Exp aenv slix
slix) (DelayedOpenAcc aenv (Array sl e) -> WithReprs (Array sl e)
forall a'. HasCallStack => DelayedOpenAcc aenv a' -> WithReprs a'
manifest DelayedOpenAcc aenv (Array sl e)
acc)
    Slice SliceIndex slix sl co sh
slice DelayedOpenAcc aenv (Array sh e)
acc Exp aenv slix
slix          -> SliceIndex slix sl co sh
-> WithReprs (Array sh e) -> slix -> WithReprs (Array sl e)
forall slix sl co sh e.
SliceIndex slix sl co sh
-> WithReprs (Array sh e) -> slix -> WithReprs (Array sl e)
sliceOp SliceIndex slix sl co sh
slice (DelayedOpenAcc aenv (Array sh e) -> WithReprs (Array sh e)
forall a'. HasCallStack => DelayedOpenAcc aenv a' -> WithReprs a'
manifest DelayedOpenAcc aenv (Array sh e)
acc) (Exp aenv slix -> slix
forall t. Exp aenv t -> t
evalE Exp aenv slix
slix)

    -- Consumers
    -- ---------
    Fold Fun aenv (e -> e -> e)
f (Just Exp aenv e
z) DelayedOpenAcc aenv (Array (sh, Int) e)
acc           -> (e -> e -> e)
-> e -> Delayed (Array (sh, Int) e) -> WithReprs (Array sh e)
forall e sh.
(e -> e -> e)
-> e -> Delayed (Array (sh, Int) e) -> WithReprs (Array sh e)
foldOp (Fun aenv (e -> e -> e) -> e -> e -> e
forall f. Fun aenv f -> f
evalF Fun aenv (e -> e -> e)
f) (Exp aenv e -> e
forall t. Exp aenv t -> t
evalE Exp aenv e
z) (DelayedOpenAcc aenv (Array (sh, Int) e)
-> Delayed (Array (sh, Int) e)
forall sh e.
DelayedOpenAcc aenv (Array sh e) -> Delayed (Array sh e)
delayed DelayedOpenAcc aenv (Array (sh, Int) e)
acc)
    Fold Fun aenv (e -> e -> e)
f Maybe (Exp aenv e)
Nothing  DelayedOpenAcc aenv (Array (sh, Int) e)
acc           -> (e -> e -> e)
-> Delayed (Array (sh, Int) e) -> WithReprs (Array sh e)
forall e sh.
HasCallStack =>
(e -> e -> e)
-> Delayed (Array (sh, Int) e) -> WithReprs (Array sh e)
fold1Op (Fun aenv (e -> e -> e) -> e -> e -> e
forall f. Fun aenv f -> f
evalF Fun aenv (e -> e -> e)
f) (DelayedOpenAcc aenv (Array (sh, Int) e)
-> Delayed (Array (sh, Int) e)
forall sh e.
DelayedOpenAcc aenv (Array sh e) -> Delayed (Array sh e)
delayed DelayedOpenAcc aenv (Array (sh, Int) e)
acc)
    FoldSeg IntegralType i
i Fun aenv (e -> e -> e)
f (Just Exp aenv e
z) DelayedOpenAcc aenv (Array (sh, Int) e)
acc DelayedOpenAcc aenv (Segments i)
seg  -> IntegralType i
-> (e -> e -> e)
-> e
-> Delayed (Array (sh, Int) e)
-> Delayed (Segments i)
-> WithReprs (Array (sh, Int) e)
forall i e sh.
HasCallStack =>
IntegralType i
-> (e -> e -> e)
-> e
-> Delayed (Array (sh, Int) e)
-> Delayed (Segments i)
-> WithReprs (Array (sh, Int) e)
foldSegOp IntegralType i
i (Fun aenv (e -> e -> e) -> e -> e -> e
forall f. Fun aenv f -> f
evalF Fun aenv (e -> e -> e)
f) (Exp aenv e -> e
forall t. Exp aenv t -> t
evalE Exp aenv e
z) (DelayedOpenAcc aenv (Array (sh, Int) e)
-> Delayed (Array (sh, Int) e)
forall sh e.
DelayedOpenAcc aenv (Array sh e) -> Delayed (Array sh e)
delayed DelayedOpenAcc aenv (Array (sh, Int) e)
acc) (DelayedOpenAcc aenv (Segments i) -> Delayed (Segments i)
forall sh e.
DelayedOpenAcc aenv (Array sh e) -> Delayed (Array sh e)
delayed DelayedOpenAcc aenv (Segments i)
seg)
    FoldSeg IntegralType i
i Fun aenv (e -> e -> e)
f Maybe (Exp aenv e)
Nothing DelayedOpenAcc aenv (Array (sh, Int) e)
acc DelayedOpenAcc aenv (Segments i)
seg   -> IntegralType i
-> (e -> e -> e)
-> Delayed (Array (sh, Int) e)
-> Delayed (Segments i)
-> WithReprs (Array (sh, Int) e)
forall i e sh.
HasCallStack =>
IntegralType i
-> (e -> e -> e)
-> Delayed (Array (sh, Int) e)
-> Delayed (Segments i)
-> WithReprs (Array (sh, Int) e)
fold1SegOp IntegralType i
i (Fun aenv (e -> e -> e) -> e -> e -> e
forall f. Fun aenv f -> f
evalF Fun aenv (e -> e -> e)
f) (DelayedOpenAcc aenv (Array (sh, Int) e)
-> Delayed (Array (sh, Int) e)
forall sh e.
DelayedOpenAcc aenv (Array sh e) -> Delayed (Array sh e)
delayed DelayedOpenAcc aenv (Array (sh, Int) e)
acc) (DelayedOpenAcc aenv (Segments i) -> Delayed (Segments i)
forall sh e.
DelayedOpenAcc aenv (Array sh e) -> Delayed (Array sh e)
delayed DelayedOpenAcc aenv (Segments i)
seg)
    Scan  Direction
d Fun aenv (e -> e -> e)
f (Just Exp aenv e
z) DelayedOpenAcc aenv (Array (sh, Int) e)
acc        -> Direction
-> ((e -> e -> e)
    -> e -> Delayed (Array (sh, Int) e) -> WithReprs a)
-> ((e -> e -> e)
    -> e -> Delayed (Array (sh, Int) e) -> WithReprs a)
-> (e -> e -> e)
-> e
-> Delayed (Array (sh, Int) e)
-> WithReprs a
forall t. Direction -> t -> t -> t
dir Direction
d (e -> e -> e) -> e -> Delayed (Array (sh, Int) e) -> WithReprs a
(e -> e -> e)
-> e
-> Delayed (Array (sh, Int) e)
-> WithReprs (Array (sh, Int) e)
forall sh e.
(e -> e -> e)
-> e
-> Delayed (Array (sh, Int) e)
-> WithReprs (Array (sh, Int) e)
scanlOp  (e -> e -> e) -> e -> Delayed (Array (sh, Int) e) -> WithReprs a
(e -> e -> e)
-> e
-> Delayed (Array (sh, Int) e)
-> WithReprs (Array (sh, Int) e)
forall sh e.
(e -> e -> e)
-> e
-> Delayed (Array (sh, Int) e)
-> WithReprs (Array (sh, Int) e)
scanrOp  (Fun aenv (e -> e -> e) -> e -> e -> e
forall f. Fun aenv f -> f
evalF Fun aenv (e -> e -> e)
f) (Exp aenv e -> e
forall t. Exp aenv t -> t
evalE Exp aenv e
z) (DelayedOpenAcc aenv (Array (sh, Int) e)
-> Delayed (Array (sh, Int) e)
forall sh e.
DelayedOpenAcc aenv (Array sh e) -> Delayed (Array sh e)
delayed DelayedOpenAcc aenv (Array (sh, Int) e)
acc)
    Scan  Direction
d Fun aenv (e -> e -> e)
f Maybe (Exp aenv e)
Nothing  DelayedOpenAcc aenv (Array (sh, Int) e)
acc        -> Direction
-> ((e -> e -> e) -> Delayed (Array (sh, Int) e) -> WithReprs a)
-> ((e -> e -> e) -> Delayed (Array (sh, Int) e) -> WithReprs a)
-> (e -> e -> e)
-> Delayed (Array (sh, Int) e)
-> WithReprs a
forall t. Direction -> t -> t -> t
dir Direction
d (e -> e -> e) -> Delayed (Array (sh, Int) e) -> WithReprs a
(e -> e -> e)
-> Delayed (Array (sh, Int) e) -> WithReprs (Array (sh, Int) e)
forall sh e.
HasCallStack =>
(e -> e -> e)
-> Delayed (Array (sh, Int) e) -> WithReprs (Array (sh, Int) e)
scanl1Op (e -> e -> e) -> Delayed (Array (sh, Int) e) -> WithReprs a
(e -> e -> e)
-> Delayed (Array (sh, Int) e) -> WithReprs (Array (sh, Int) e)
forall sh e.
HasCallStack =>
(e -> e -> e)
-> Delayed (Array (sh, Int) e) -> WithReprs (Array (sh, Int) e)
scanr1Op (Fun aenv (e -> e -> e) -> e -> e -> e
forall f. Fun aenv f -> f
evalF Fun aenv (e -> e -> e)
f)           (DelayedOpenAcc aenv (Array (sh, Int) e)
-> Delayed (Array (sh, Int) e)
forall sh e.
DelayedOpenAcc aenv (Array sh e) -> Delayed (Array sh e)
delayed DelayedOpenAcc aenv (Array (sh, Int) e)
acc)
    Scan' Direction
d Fun aenv (e -> e -> e)
f Exp aenv e
z DelayedOpenAcc aenv (Array (sh, Int) e)
acc               -> Direction
-> ((e -> e -> e)
    -> e -> Delayed (Array (sh, Int) e) -> WithReprs a)
-> ((e -> e -> e)
    -> e -> Delayed (Array (sh, Int) e) -> WithReprs a)
-> (e -> e -> e)
-> e
-> Delayed (Array (sh, Int) e)
-> WithReprs a
forall t. Direction -> t -> t -> t
dir Direction
d (e -> e -> e) -> e -> Delayed (Array (sh, Int) e) -> WithReprs a
(e -> e -> e)
-> e
-> Delayed (Array (sh, Int) e)
-> WithReprs (Array (sh, Int) e, Array sh e)
forall sh e.
(e -> e -> e)
-> e
-> Delayed (Array (sh, Int) e)
-> WithReprs (Array (sh, Int) e, Array sh e)
scanl'Op (e -> e -> e) -> e -> Delayed (Array (sh, Int) e) -> WithReprs a
(e -> e -> e)
-> e
-> Delayed (Array (sh, Int) e)
-> WithReprs (Array (sh, Int) e, Array sh e)
forall sh e.
(e -> e -> e)
-> e
-> Delayed (Array (sh, Int) e)
-> WithReprs (Array (sh, Int) e, Array sh e)
scanr'Op (Fun aenv (e -> e -> e) -> e -> e -> e
forall f. Fun aenv f -> f
evalF Fun aenv (e -> e -> e)
f) (Exp aenv e -> e
forall t. Exp aenv t -> t
evalE Exp aenv e
z) (DelayedOpenAcc aenv (Array (sh, Int) e)
-> Delayed (Array (sh, Int) e)
forall sh e.
DelayedOpenAcc aenv (Array sh e) -> Delayed (Array sh e)
delayed DelayedOpenAcc aenv (Array (sh, Int) e)
acc)
    Permute Fun aenv (e -> e -> e)
f DelayedOpenAcc aenv (Array sh' e)
def Fun aenv (sh -> PrimMaybe sh')
p DelayedOpenAcc aenv (Array sh e)
acc           -> (e -> e -> e)
-> WithReprs (Array sh' e)
-> (sh -> PrimMaybe sh')
-> Delayed (Array sh e)
-> WithReprs (Array sh' e)
forall sh sh' e.
HasCallStack =>
(e -> e -> e)
-> WithReprs (Array sh' e)
-> (sh -> PrimMaybe sh')
-> Delayed (Array sh e)
-> WithReprs (Array sh' e)
permuteOp (Fun aenv (e -> e -> e) -> e -> e -> e
forall f. Fun aenv f -> f
evalF Fun aenv (e -> e -> e)
f) (DelayedOpenAcc aenv (Array sh' e) -> WithReprs (Array sh' e)
forall a'. HasCallStack => DelayedOpenAcc aenv a' -> WithReprs a'
manifest DelayedOpenAcc aenv (Array sh' e)
def) (Fun aenv (sh -> PrimMaybe sh') -> sh -> PrimMaybe sh'
forall f. Fun aenv f -> f
evalF Fun aenv (sh -> PrimMaybe sh')
p) (DelayedOpenAcc aenv (Array sh e) -> Delayed (Array sh e)
forall sh e.
DelayedOpenAcc aenv (Array sh e) -> Delayed (Array sh e)
delayed DelayedOpenAcc aenv (Array sh e)
acc)
    Stencil StencilR sh e stencil
s TypeR e'
tp Fun aenv (stencil -> e')
sten Boundary aenv (Array sh e)
b DelayedOpenAcc aenv (Array sh e)
acc       -> StencilR sh e stencil
-> TypeR e'
-> (stencil -> e')
-> Boundary (Array sh e)
-> Delayed (Array sh e)
-> WithReprs (Array sh e')
forall sh a stencil b.
HasCallStack =>
StencilR sh a stencil
-> TypeR b
-> (stencil -> b)
-> Boundary (Array sh a)
-> Delayed (Array sh a)
-> WithReprs (Array sh b)
stencilOp StencilR sh e stencil
s TypeR e'
tp (Fun aenv (stencil -> e') -> stencil -> e'
forall f. Fun aenv f -> f
evalF Fun aenv (stencil -> e')
sten) (Boundary aenv (Array sh e) -> Boundary (Array sh e)
forall t. Boundary aenv t -> Boundary t
evalB Boundary aenv (Array sh e)
b) (DelayedOpenAcc aenv (Array sh e) -> Delayed (Array sh e)
forall sh e.
DelayedOpenAcc aenv (Array sh e) -> Delayed (Array sh e)
delayed DelayedOpenAcc aenv (Array sh e)
acc)
    Stencil2 StencilR sh a1 stencil1
s1 StencilR sh b stencil2
s2 TypeR c
tp Fun aenv (stencil1 -> stencil2 -> c)
sten Boundary aenv (Array sh a1)
b1 DelayedOpenAcc aenv (Array sh a1)
a1 Boundary aenv (Array sh b)
b2 DelayedOpenAcc aenv (Array sh b)
a2
                                  -> StencilR sh a1 stencil1
-> StencilR sh b stencil2
-> TypeR c
-> (stencil1 -> stencil2 -> c)
-> Boundary (Array sh a1)
-> Delayed (Array sh a1)
-> Boundary (Array sh b)
-> Delayed (Array sh b)
-> WithReprs (Array sh c)
forall sh a stencil1 b stencil2 c.
HasCallStack =>
StencilR sh a stencil1
-> StencilR sh b stencil2
-> TypeR c
-> (stencil1 -> stencil2 -> c)
-> Boundary (Array sh a)
-> Delayed (Array sh a)
-> Boundary (Array sh b)
-> Delayed (Array sh b)
-> WithReprs (Array sh c)
stencil2Op StencilR sh a1 stencil1
s1 StencilR sh b stencil2
s2 TypeR c
tp (Fun aenv (stencil1 -> stencil2 -> c) -> stencil1 -> stencil2 -> c
forall f. Fun aenv f -> f
evalF Fun aenv (stencil1 -> stencil2 -> c)
sten) (Boundary aenv (Array sh a1) -> Boundary (Array sh a1)
forall t. Boundary aenv t -> Boundary t
evalB Boundary aenv (Array sh a1)
b1) (DelayedOpenAcc aenv (Array sh a1) -> Delayed (Array sh a1)
forall sh e.
DelayedOpenAcc aenv (Array sh e) -> Delayed (Array sh e)
delayed DelayedOpenAcc aenv (Array sh a1)
a1) (Boundary aenv (Array sh b) -> Boundary (Array sh b)
forall t. Boundary aenv t -> Boundary t
evalB Boundary aenv (Array sh b)
b2) (DelayedOpenAcc aenv (Array sh b) -> Delayed (Array sh b)
forall sh e.
DelayedOpenAcc aenv (Array sh e) -> Delayed (Array sh e)
delayed DelayedOpenAcc aenv (Array sh b)
a2)


-- Array primitives
-- ----------------

unitOp :: TypeR e -> e -> WithReprs (Scalar e)
unitOp :: forall e. TypeR e -> e -> WithReprs (Scalar e)
unitOp TypeR e
tp e
e = ArrayR (Array () e) -> () -> (() -> e) -> WithReprs (Array () e)
forall sh e.
ArrayR (Array sh e) -> sh -> (sh -> e) -> WithReprs (Array sh e)
fromFunction' (ShapeR () -> TypeR e -> ArrayR (Array () e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR ()
ShapeRz TypeR e
tp) () (e -> () -> e
forall a b. a -> b -> a
const e
e)


generateOp
    :: ArrayR (Array sh e)
    -> sh
    -> (sh -> e)
    -> WithReprs (Array sh e)
generateOp :: forall sh e.
ArrayR (Array sh e) -> sh -> (sh -> e) -> WithReprs (Array sh e)
generateOp = ArrayR (Array sh e) -> sh -> (sh -> e) -> WithReprs (Array sh e)
forall sh e.
ArrayR (Array sh e) -> sh -> (sh -> e) -> WithReprs (Array sh e)
fromFunction'


transformOp
    :: ArrayR (Array sh' b)
    -> sh'
    -> (sh' -> sh)
    -> (a -> b)
    -> Delayed (Array sh a)
    -> WithReprs (Array sh' b)
transformOp :: forall sh' b sh a.
ArrayR (Array sh' b)
-> sh'
-> (sh' -> sh)
-> (a -> b)
-> Delayed (Array sh a)
-> WithReprs (Array sh' b)
transformOp ArrayR (Array sh' b)
repr sh'
sh' sh' -> sh
p a -> b
f (Delayed ArrayR (Array sh e)
_ sh
_ sh -> e
xs Int -> e
_)
  = ArrayR (Array sh' b)
-> sh' -> (sh' -> b) -> WithReprs (Array sh' b)
forall sh e.
ArrayR (Array sh e) -> sh -> (sh -> e) -> WithReprs (Array sh e)
fromFunction' ArrayR (Array sh' b)
repr sh'
sh' (\sh'
ix -> a -> b
f (sh -> e
xs (sh -> e) -> sh -> e
forall a b. (a -> b) -> a -> b
$ sh' -> sh
p sh'
ix))


reshapeOp
    :: HasCallStack
    => ShapeR sh
    -> sh
    -> WithReprs (Array sh' e)
    -> WithReprs (Array sh  e)
reshapeOp :: forall sh sh' e.
HasCallStack =>
ShapeR sh
-> sh -> WithReprs (Array sh' e) -> WithReprs (Array sh e)
reshapeOp ShapeR sh
newShapeR sh
newShape (TupRsingle (ArrayR ShapeR sh
shr TypeR e
tp), (Array sh'
sh ArrayData e
adata))
  = Builder -> Bool -> WithReprs (Array sh e) -> WithReprs (Array sh e)
forall a. HasCallStack => Builder -> Bool -> a -> a
boundsCheck Builder
"shape mismatch" (ShapeR sh -> sh -> Int
forall sh. ShapeR sh -> sh -> Int
size ShapeR sh
newShapeR sh
newShape Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== ShapeR sh -> sh -> Int
forall sh. ShapeR sh -> sh -> Int
size ShapeR sh
shr sh'
sh
sh)
    ( ArrayR (Array sh e) -> TupR ArrayR (Array sh e)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ShapeR sh -> TypeR e -> ArrayR (Array sh e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh
newShapeR TypeR e
TypeR e
tp)
    , sh -> ArrayData e -> Array sh e
forall sh e. sh -> ArrayData e -> Array sh e
Array sh
newShape ArrayData e
adata
    )


replicateOp
    :: SliceIndex slix sl co sh
    -> slix
    -> WithReprs (Array sl e)
    -> WithReprs (Array sh e)
replicateOp :: forall slix sl co sh e.
SliceIndex slix sl co sh
-> slix -> WithReprs (Array sl e) -> WithReprs (Array sh e)
replicateOp SliceIndex slix sl co sh
slice slix
slix (TupRsingle repr :: ArrayR (Array sl e)
repr@(ArrayR ShapeR sh
_ TypeR e
tp), Array sl e
arr)
  = ArrayR (Array sh e) -> sh -> (sh -> e) -> WithReprs (Array sh e)
forall sh e.
ArrayR (Array sh e) -> sh -> (sh -> e) -> WithReprs (Array sh e)
fromFunction' ArrayR (Array sh e)
ArrayR (Array sh e)
repr' sh
sh (\sh
ix -> (ArrayR (Array sl e)
repr, Array sl e
arr) (ArrayR (Array sl e), Array sl e) -> sl -> e
forall sh e. (ArrayR (Array sh e), Array sh e) -> sh -> e
! sh -> sl
pf sh
ix)
  where
    repr' :: ArrayR (Array sh e)
repr' = ShapeR sh -> TypeR e -> ArrayR (Array sh e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR (SliceIndex slix sl co sh -> ShapeR sh
forall slix sl co dim. SliceIndex slix sl co dim -> ShapeR dim
sliceDomainR SliceIndex slix sl co sh
slice) TypeR e
tp
    (sh
sh, sh -> sl
pf) = SliceIndex slix sl co sh -> slix -> sl -> (sh, sh -> sl)
forall slix sl co dim.
SliceIndex slix sl co dim -> slix -> sl -> (dim, dim -> sl)
extend SliceIndex slix sl co sh
slice slix
slix (Array sl e -> sl
forall sh e. Array sh e -> sh
shape Array sl e
arr)

    extend :: SliceIndex slix sl co dim
           -> slix
           -> sl
           -> (dim, dim -> sl)
    extend :: forall slix sl co dim.
SliceIndex slix sl co dim -> slix -> sl -> (dim, dim -> sl)
extend SliceIndex slix sl co dim
SliceNil              ()        ()
      = ((), sl -> dim -> sl
forall a b. a -> b -> a
const ())
    extend (SliceAll SliceIndex ix1 slice1 co dim
sliceIdx)   (ix1
slx, ()) (slice1
sl, Int
sz)
      = let (dim
dim', dim -> slice1
f') = SliceIndex ix1 slice1 co dim
-> ix1 -> slice1 -> (dim, dim -> slice1)
forall slix sl co dim.
SliceIndex slix sl co dim -> slix -> sl -> (dim, dim -> sl)
extend SliceIndex ix1 slice1 co dim
sliceIdx ix1
slx slice1
sl
        in  ((dim
dim', Int
sz), \(dim
ix, Int
i) -> (dim -> slice1
f' dim
ix, Int
i))
    extend (SliceFixed SliceIndex ix1 sl co dim
sliceIdx) (ix1
slx, Int
sz) sl
sl
      = let (dim
dim', dim -> sl
f') = SliceIndex ix1 sl co dim -> ix1 -> sl -> (dim, dim -> sl)
forall slix sl co dim.
SliceIndex slix sl co dim -> slix -> sl -> (dim, dim -> sl)
extend SliceIndex ix1 sl co dim
sliceIdx ix1
slx sl
sl
        in  ((dim
dim', Int
sz), \(dim
ix, Int
_) -> dim -> sl
f' dim
ix)


sliceOp
    :: SliceIndex slix sl co sh
    -> WithReprs (Array sh e)
    -> slix
    -> WithReprs (Array sl e)
sliceOp :: forall slix sl co sh e.
SliceIndex slix sl co sh
-> WithReprs (Array sh e) -> slix -> WithReprs (Array sl e)
sliceOp SliceIndex slix sl co sh
slice (TupRsingle repr :: ArrayR (Array sh e)
repr@(ArrayR ShapeR sh
_ TypeR e
tp), Array sh e
arr) slix
slix
  = ArrayR (Array sl e) -> sl -> (sl -> e) -> WithReprs (Array sl e)
forall sh e.
ArrayR (Array sh e) -> sh -> (sh -> e) -> WithReprs (Array sh e)
fromFunction' ArrayR (Array sl e)
ArrayR (Array sl e)
repr' sl
sh' (\sl
ix -> (ArrayR (Array sh e)
repr, Array sh e
arr) (ArrayR (Array sh e), Array sh e) -> sh -> e
forall sh e. (ArrayR (Array sh e), Array sh e) -> sh -> e
! sl -> sh
pf sl
ix)
  where
    repr' :: ArrayR (Array sl e)
repr' = ShapeR sl -> TypeR e -> ArrayR (Array sl e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR (SliceIndex slix sl co sh -> ShapeR sl
forall slix sl co dim. SliceIndex slix sl co dim -> ShapeR sl
sliceShapeR SliceIndex slix sl co sh
slice) TypeR e
tp
    (sl
sh', sl -> sh
pf) = SliceIndex slix sl co sh -> slix -> sh -> (sl, sl -> sh)
forall slix sl co sh.
HasCallStack =>
SliceIndex slix sl co sh -> slix -> sh -> (sl, sl -> sh)
restrict SliceIndex slix sl co sh
slice slix
slix (Array sh e -> sh
forall sh e. Array sh e -> sh
shape Array sh e
arr)

    restrict
        :: HasCallStack
        => SliceIndex slix sl co sh
        -> slix
        -> sh
        -> (sl, sl -> sh)
    restrict :: forall slix sl co sh.
HasCallStack =>
SliceIndex slix sl co sh -> slix -> sh -> (sl, sl -> sh)
restrict SliceIndex slix sl co sh
SliceNil              ()        ()
      = ((), sh -> sl -> sh
forall a b. a -> b -> a
const ())
    restrict (SliceAll SliceIndex ix1 slice1 co dim
sliceIdx)   (ix1
slx, ()) (dim
sl, Int
sz)
      = let (slice1
sl', slice1 -> dim
f') = SliceIndex ix1 slice1 co dim
-> ix1 -> dim -> (slice1, slice1 -> dim)
forall slix sl co sh.
HasCallStack =>
SliceIndex slix sl co sh -> slix -> sh -> (sl, sl -> sh)
restrict SliceIndex ix1 slice1 co dim
sliceIdx ix1
slx dim
sl
        in  ((slice1
sl', Int
sz), \(slice1
ix, Int
i) -> (slice1 -> dim
f' slice1
ix, Int
i))
    restrict (SliceFixed SliceIndex ix1 sl co dim
sliceIdx) (ix1
slx, Int
i)  (dim
sl, Int
sz)
      = let (sl
sl', sl -> dim
f') = SliceIndex ix1 sl co dim -> ix1 -> dim -> (sl, sl -> dim)
forall slix sl co sh.
HasCallStack =>
SliceIndex slix sl co sh -> slix -> sh -> (sl, sl -> sh)
restrict SliceIndex ix1 sl co dim
sliceIdx ix1
slx dim
sl
        in  Int -> Int -> (sl, sl -> sh) -> (sl, sl -> sh)
forall a. HasCallStack => Int -> Int -> a -> a
indexCheck Int
i Int
sz ((sl, sl -> sh) -> (sl, sl -> sh))
-> (sl, sl -> sh) -> (sl, sl -> sh)
forall a b. (a -> b) -> a -> b
$ (sl
sl', \sl
ix -> (sl -> dim
f' sl
ix, Int
i))


mapOp :: TypeR b
      -> (a -> b)
      -> Delayed   (Array sh a)
      -> WithReprs (Array sh b)
mapOp :: forall b a sh.
TypeR b
-> (a -> b) -> Delayed (Array sh a) -> WithReprs (Array sh b)
mapOp TypeR b
tp a -> b
f (Delayed (ArrayR ShapeR sh
shr TypeR e
_) sh
sh sh -> e
xs Int -> e
_)
  = ArrayR (Array sh b) -> sh -> (sh -> b) -> WithReprs (Array sh b)
forall sh e.
ArrayR (Array sh e) -> sh -> (sh -> e) -> WithReprs (Array sh e)
fromFunction' (ShapeR sh -> TypeR b -> ArrayR (Array sh b)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh
ShapeR sh
shr TypeR b
tp) sh
sh
sh (\sh
ix -> a -> b
f (sh -> e
xs sh
sh
ix))


zipWithOp
    :: TypeR c
    -> (a -> b -> c)
    -> Delayed   (Array sh a)
    -> Delayed   (Array sh b)
    -> WithReprs (Array sh c)
zipWithOp :: forall c a b sh.
TypeR c
-> (a -> b -> c)
-> Delayed (Array sh a)
-> Delayed (Array sh b)
-> WithReprs (Array sh c)
zipWithOp TypeR c
tp a -> b -> c
f (Delayed (ArrayR ShapeR sh
shr TypeR e
_) sh
shx sh -> e
xs Int -> e
_) (Delayed ArrayR (Array sh e)
_ sh
shy sh -> e
ys Int -> e
_)
  = ArrayR (Array sh c) -> sh -> (sh -> c) -> WithReprs (Array sh c)
forall sh e.
ArrayR (Array sh e) -> sh -> (sh -> e) -> WithReprs (Array sh e)
fromFunction' (ShapeR sh -> TypeR c -> ArrayR (Array sh c)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh
ShapeR sh
shr TypeR c
tp) (ShapeR sh -> sh -> sh -> sh
forall sh. ShapeR sh -> sh -> sh -> sh
intersect ShapeR sh
ShapeR sh
shr sh
sh
shx sh
sh
shy) (\sh
ix -> a -> b -> c
f (sh -> e
xs sh
sh
ix) (sh -> e
ys sh
sh
ix))


foldOp
    :: (e -> e -> e)
    -> e
    -> Delayed (Array (sh, Int) e)
    -> WithReprs (Array sh e)
foldOp :: forall e sh.
(e -> e -> e)
-> e -> Delayed (Array (sh, Int) e) -> WithReprs (Array sh e)
foldOp e -> e -> e
f e
z (Delayed (ArrayR (ShapeRsnoc ShapeR sh1
shr) TypeR e
tp) (sh
sh, Int
n) sh -> e
arr Int -> e
_)
  = ArrayR (Array sh e) -> sh -> (sh -> e) -> WithReprs (Array sh e)
forall sh e.
ArrayR (Array sh e) -> sh -> (sh -> e) -> WithReprs (Array sh e)
fromFunction' (ShapeR sh -> TypeR e -> ArrayR (Array sh e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh
ShapeR sh1
shr TypeR e
TypeR e
tp) sh
sh (\sh
ix -> ShapeR ((), Int)
-> ((), Int) -> (((), Int) -> e) -> (e -> e -> e) -> e -> e
forall sh a.
ShapeR sh -> sh -> (sh -> a) -> (a -> a -> a) -> a -> a
iter (ShapeR () -> ShapeR ((), Int)
forall sh1. ShapeR sh1 -> ShapeR (sh1, Int)
ShapeRsnoc ShapeR ()
ShapeRz) ((), Int
n) (\((), Int
i) -> sh -> e
arr (sh
ix, Int
i)) e -> e -> e
f e
z)


fold1Op
    :: HasCallStack
    => (e -> e -> e)
    -> Delayed (Array (sh, Int) e)
    -> WithReprs (Array sh e)
fold1Op :: forall e sh.
HasCallStack =>
(e -> e -> e)
-> Delayed (Array (sh, Int) e) -> WithReprs (Array sh e)
fold1Op e -> e -> e
f (Delayed (ArrayR (ShapeRsnoc ShapeR sh1
shr) TypeR e
tp) (sh
sh, Int
n) sh -> e
arr Int -> e
_)
  = Builder -> Bool -> WithReprs (Array sh e) -> WithReprs (Array sh e)
forall a. HasCallStack => Builder -> Bool -> a -> a
boundsCheck Builder
"empty array" (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0)
  (WithReprs (Array sh e) -> WithReprs (Array sh e))
-> WithReprs (Array sh e) -> WithReprs (Array sh e)
forall a b. (a -> b) -> a -> b
$ ArrayR (Array sh e) -> sh -> (sh -> e) -> WithReprs (Array sh e)
forall sh e.
ArrayR (Array sh e) -> sh -> (sh -> e) -> WithReprs (Array sh e)
fromFunction' (ShapeR sh -> TypeR e -> ArrayR (Array sh e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh
ShapeR sh1
shr TypeR e
TypeR e
tp) sh
sh (\sh
ix -> ShapeR ((), Int)
-> ((), Int) -> (((), Int) -> e) -> (e -> e -> e) -> e
forall sh a.
HasCallStack =>
ShapeR sh -> sh -> (sh -> a) -> (a -> a -> a) -> a
iter1 (ShapeR () -> ShapeR ((), Int)
forall sh1. ShapeR sh1 -> ShapeR (sh1, Int)
ShapeRsnoc ShapeR ()
ShapeRz) ((), Int
n) (\((), Int
i) -> sh -> e
arr (sh
ix, Int
i)) e -> e -> e
f)


foldSegOp
    :: HasCallStack
    => IntegralType i
    -> (e -> e -> e)
    -> e
    -> Delayed (Array (sh, Int) e)
    -> Delayed (Segments i)
    -> WithReprs (Array (sh, Int) e)
foldSegOp :: forall i e sh.
HasCallStack =>
IntegralType i
-> (e -> e -> e)
-> e
-> Delayed (Array (sh, Int) e)
-> Delayed (Segments i)
-> WithReprs (Array (sh, Int) e)
foldSegOp IntegralType i
itp e -> e -> e
f e
z (Delayed ArrayR (Array sh e)
repr (sh
sh, Int
_) sh -> e
arr Int -> e
_) (Delayed ArrayR (Array sh e)
_ ((), Int
n) sh -> e
_ Int -> e
seg)
  | IntegralDict i
IntegralDict <- IntegralType i -> IntegralDict i
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType i
itp
  = Builder
-> Bool
-> WithReprs (Array (sh, Int) e)
-> WithReprs (Array (sh, Int) e)
forall a. HasCallStack => Builder -> Bool -> a -> a
boundsCheck Builder
"empty segment descriptor" (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0)
  (WithReprs (Array (sh, Int) e) -> WithReprs (Array (sh, Int) e))
-> WithReprs (Array (sh, Int) e) -> WithReprs (Array (sh, Int) e)
forall a b. (a -> b) -> a -> b
$ ArrayR (Array (sh, Int) e)
-> (sh, Int) -> ((sh, Int) -> e) -> WithReprs (Array (sh, Int) e)
forall sh e.
ArrayR (Array sh e) -> sh -> (sh -> e) -> WithReprs (Array sh e)
fromFunction' ArrayR (Array sh e)
ArrayR (Array (sh, Int) e)
repr (sh
sh, Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
  (((sh, Int) -> e) -> WithReprs (Array (sh, Int) e))
-> ((sh, Int) -> e) -> WithReprs (Array (sh, Int) e)
forall a b. (a -> b) -> a -> b
$ \(sh
sz, Int
ix) -> let start :: Int
start = e -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (e -> Int) -> e -> Int
forall a b. (a -> b) -> a -> b
$ Int -> e
seg Int
ix
                     end :: Int
end   = e -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (e -> Int) -> e -> Int
forall a b. (a -> b) -> a -> b
$ Int -> e
seg (Int
ixInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
                 in
                 Builder -> Bool -> e -> e
forall a. HasCallStack => Builder -> Bool -> a -> a
boundsCheck Builder
"empty segment" (Int
end Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
start)
                 (e -> e) -> e -> e
forall a b. (a -> b) -> a -> b
$ ShapeR ((), Int)
-> ((), Int) -> (((), Int) -> e) -> (e -> e -> e) -> e -> e
forall sh a.
ShapeR sh -> sh -> (sh -> a) -> (a -> a -> a) -> a -> a
iter (ShapeR () -> ShapeR ((), Int)
forall sh1. ShapeR sh1 -> ShapeR (sh1, Int)
ShapeRsnoc ShapeR ()
ShapeRz) ((), Int
endInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
start) (\((), Int
i) -> sh -> e
arr (sh
sz, Int
startInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
i)) e -> e -> e
f e
z


fold1SegOp
    :: HasCallStack
    => IntegralType i
    -> (e -> e -> e)
    -> Delayed (Array (sh, Int) e)
    -> Delayed (Segments i)
    -> WithReprs (Array (sh, Int) e)
fold1SegOp :: forall i e sh.
HasCallStack =>
IntegralType i
-> (e -> e -> e)
-> Delayed (Array (sh, Int) e)
-> Delayed (Segments i)
-> WithReprs (Array (sh, Int) e)
fold1SegOp IntegralType i
itp e -> e -> e
f (Delayed ArrayR (Array sh e)
repr (sh
sh, Int
_) sh -> e
arr Int -> e
_) (Delayed ArrayR (Array sh e)
_ ((), Int
n) sh -> e
_ Int -> e
seg)
  | IntegralDict i
IntegralDict <- IntegralType i -> IntegralDict i
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType i
itp
  = Builder
-> Bool
-> WithReprs (Array (sh, Int) e)
-> WithReprs (Array (sh, Int) e)
forall a. HasCallStack => Builder -> Bool -> a -> a
boundsCheck Builder
"empty segment descriptor" (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0)
  (WithReprs (Array (sh, Int) e) -> WithReprs (Array (sh, Int) e))
-> WithReprs (Array (sh, Int) e) -> WithReprs (Array (sh, Int) e)
forall a b. (a -> b) -> a -> b
$ ArrayR (Array (sh, Int) e)
-> (sh, Int) -> ((sh, Int) -> e) -> WithReprs (Array (sh, Int) e)
forall sh e.
ArrayR (Array sh e) -> sh -> (sh -> e) -> WithReprs (Array sh e)
fromFunction' ArrayR (Array sh e)
ArrayR (Array (sh, Int) e)
repr (sh
sh, Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
  (((sh, Int) -> e) -> WithReprs (Array (sh, Int) e))
-> ((sh, Int) -> e) -> WithReprs (Array (sh, Int) e)
forall a b. (a -> b) -> a -> b
$ \(sh
sz, Int
ix)   -> let start :: Int
start = e -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (e -> Int) -> e -> Int
forall a b. (a -> b) -> a -> b
$ Int -> e
seg Int
ix
                       end :: Int
end   = e -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (e -> Int) -> e -> Int
forall a b. (a -> b) -> a -> b
$ Int -> e
seg (Int
ixInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
                   in
                   Builder -> Bool -> e -> e
forall a. HasCallStack => Builder -> Bool -> a -> a
boundsCheck Builder
"empty segment" (Int
end Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
start)
                   (e -> e) -> e -> e
forall a b. (a -> b) -> a -> b
$ ShapeR ((), Int)
-> ((), Int) -> (((), Int) -> e) -> (e -> e -> e) -> e
forall sh a.
HasCallStack =>
ShapeR sh -> sh -> (sh -> a) -> (a -> a -> a) -> a
iter1 (ShapeR () -> ShapeR ((), Int)
forall sh1. ShapeR sh1 -> ShapeR (sh1, Int)
ShapeRsnoc ShapeR ()
ShapeRz) ((), Int
endInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
start) (\((), Int
i) -> sh -> e
arr (sh
sz, Int
startInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
i)) e -> e -> e
f


scanl1Op
    :: forall sh e. HasCallStack
    => (e -> e -> e)
    -> Delayed (Array (sh, Int) e)
    -> WithReprs (Array (sh, Int) e)
scanl1Op :: forall sh e.
HasCallStack =>
(e -> e -> e)
-> Delayed (Array (sh, Int) e) -> WithReprs (Array (sh, Int) e)
scanl1Op e -> e -> e
f (Delayed (ArrayR ShapeR sh
shr TypeR e
tp) sh
sh sh -> e
ain Int -> e
_)
  = ( ArrayR (Array (sh, Int) e) -> TupR ArrayR (Array (sh, Int) e)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ArrayR (Array (sh, Int) e) -> TupR ArrayR (Array (sh, Int) e))
-> ArrayR (Array (sh, Int) e) -> TupR ArrayR (Array (sh, Int) e)
forall a b. (a -> b) -> a -> b
$ ShapeR (sh, Int) -> TypeR e -> ArrayR (Array (sh, Int) e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh
ShapeR (sh, Int)
shr TypeR e
TypeR e
tp
    , ArrayData e
adata ArrayData e -> Array (sh, Int) e -> Array (sh, Int) e
forall a b. a -> b -> b
`seq` (sh, Int) -> ArrayData e -> Array (sh, Int) e
forall sh e. sh -> ArrayData e -> Array sh e
Array sh
(sh, Int)
sh ArrayData e
adata
    )
  where
    --
    (ArrayData e
adata, e
_)  = forall e. IO (MutableArrayData e, e) -> (MutableArrayData e, e)
runArrayData @e (IO (ArrayData e, e) -> (ArrayData e, e))
-> IO (ArrayData e, e) -> (ArrayData e, e)
forall a b. (a -> b) -> a -> b
$ do
      ArrayData e
aout <- TypeR e -> Int -> IO (MutableArrayData e)
forall e.
HasCallStack =>
TupR ScalarType e -> Int -> IO (MutableArrayData e)
newArrayData TypeR e
tp (ShapeR sh -> sh -> Int
forall sh. ShapeR sh -> sh -> Int
size ShapeR sh
shr sh
sh
sh)

      let write :: (sh, Int) -> IO ()
write (sh
sz, Int
0) = TypeR e -> MutableArrayData e -> Int -> e -> IO ()
forall e.
TupR ScalarType e -> MutableArrayData e -> Int -> e -> IO ()
writeArrayData TypeR e
tp ArrayData e
MutableArrayData e
aout (ShapeR sh -> sh -> sh -> Int
forall sh. HasCallStack => ShapeR sh -> sh -> sh -> Int
toIndex ShapeR sh
shr sh
sh
sh (sh
sz, Int
0)) (sh -> e
ain (sh
sz, Int
0))
          write (sh
sz, Int
i) = do
            e
x <- TypeR e -> MutableArrayData e -> Int -> IO e
forall e. TupR ScalarType e -> MutableArrayData e -> Int -> IO e
readArrayData TypeR e
tp ArrayData e
MutableArrayData e
aout (ShapeR sh -> sh -> sh -> Int
forall sh. HasCallStack => ShapeR sh -> sh -> sh -> Int
toIndex ShapeR sh
shr sh
sh
sh (sh
sz, Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1))
            let y :: e
y = sh -> e
ain (sh
sz, Int
i)
            TypeR e -> MutableArrayData e -> Int -> e -> IO ()
forall e.
TupR ScalarType e -> MutableArrayData e -> Int -> e -> IO ()
writeArrayData TypeR e
tp ArrayData e
MutableArrayData e
aout (ShapeR sh -> sh -> sh -> Int
forall sh. HasCallStack => ShapeR sh -> sh -> sh -> Int
toIndex ShapeR sh
shr sh
sh
sh (sh
sz, Int
i)) (e -> e -> e
f e
e
x e
e
y)

      ShapeR sh
-> sh
-> (sh -> IO ())
-> (IO () -> IO () -> IO ())
-> IO ()
-> IO ()
forall sh a.
ShapeR sh -> sh -> (sh -> a) -> (a -> a -> a) -> a -> a
iter ShapeR sh
shr sh
sh
sh sh -> IO ()
(sh, Int) -> IO ()
write IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
(>>) (() -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ())
      (ArrayData e, e) -> IO (ArrayData e, e)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ArrayData e
aout, e
forall a. HasCallStack => a
undefined)


scanlOp
    :: forall sh e.
       (e -> e -> e)
    -> e
    -> Delayed (Array (sh, Int) e)
    -> WithReprs (Array (sh, Int) e)
scanlOp :: forall sh e.
(e -> e -> e)
-> e
-> Delayed (Array (sh, Int) e)
-> WithReprs (Array (sh, Int) e)
scanlOp e -> e -> e
f e
z (Delayed (ArrayR ShapeR sh
shr TypeR e
tp) (sh
sh, Int
n) sh -> e
ain Int -> e
_)
  = ( ArrayR (Array (sh, Int) e) -> TupR ArrayR (Array (sh, Int) e)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ArrayR (Array (sh, Int) e) -> TupR ArrayR (Array (sh, Int) e))
-> ArrayR (Array (sh, Int) e) -> TupR ArrayR (Array (sh, Int) e)
forall a b. (a -> b) -> a -> b
$ ShapeR (sh, Int) -> TypeR e -> ArrayR (Array (sh, Int) e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh
ShapeR (sh, Int)
shr TypeR e
TypeR e
tp
    , ArrayData e
adata ArrayData e -> Array (sh, Int) e -> Array (sh, Int) e
forall a b. a -> b -> b
`seq` (sh, Int) -> ArrayData e -> Array (sh, Int) e
forall sh e. sh -> ArrayData e -> Array sh e
Array (sh, Int)
sh' ArrayData e
adata
    )
  where
    sh' :: (sh, Int)
sh'         = (sh
sh, Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
    --
    (ArrayData e
adata, e
_)  = forall e. IO (MutableArrayData e, e) -> (MutableArrayData e, e)
runArrayData @e (IO (ArrayData e, e) -> (ArrayData e, e))
-> IO (ArrayData e, e) -> (ArrayData e, e)
forall a b. (a -> b) -> a -> b
$ do
      ArrayData e
aout <- TypeR e -> Int -> IO (MutableArrayData e)
forall e.
HasCallStack =>
TupR ScalarType e -> Int -> IO (MutableArrayData e)
newArrayData TypeR e
tp (ShapeR sh -> sh -> Int
forall sh. ShapeR sh -> sh -> Int
size ShapeR sh
shr sh
(sh, Int)
sh')

      let write :: (sh, Int) -> IO ()
write (sh
sz, Int
0) = TypeR e -> MutableArrayData e -> Int -> e -> IO ()
forall e.
TupR ScalarType e -> MutableArrayData e -> Int -> e -> IO ()
writeArrayData TypeR e
tp ArrayData e
MutableArrayData e
aout (ShapeR sh -> sh -> sh -> Int
forall sh. HasCallStack => ShapeR sh -> sh -> sh -> Int
toIndex ShapeR sh
shr sh
(sh, Int)
sh' (sh
sz, Int
0)) e
e
z
          write (sh
sz, Int
i) = do
            e
x <- TypeR e -> MutableArrayData e -> Int -> IO e
forall e. TupR ScalarType e -> MutableArrayData e -> Int -> IO e
readArrayData TypeR e
tp ArrayData e
MutableArrayData e
aout (ShapeR sh -> sh -> sh -> Int
forall sh. HasCallStack => ShapeR sh -> sh -> sh -> Int
toIndex ShapeR sh
shr sh
(sh, Int)
sh' (sh
sz, Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1))
            let y :: e
y = sh -> e
ain (sh
sz, Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
            TypeR e -> MutableArrayData e -> Int -> e -> IO ()
forall e.
TupR ScalarType e -> MutableArrayData e -> Int -> e -> IO ()
writeArrayData TypeR e
tp ArrayData e
MutableArrayData e
aout (ShapeR sh -> sh -> sh -> Int
forall sh. HasCallStack => ShapeR sh -> sh -> sh -> Int
toIndex ShapeR sh
shr sh
(sh, Int)
sh' (sh
sz, Int
i)) (e -> e -> e
f e
e
x e
e
y)

      ShapeR sh
-> sh
-> (sh -> IO ())
-> (IO () -> IO () -> IO ())
-> IO ()
-> IO ()
forall sh a.
ShapeR sh -> sh -> (sh -> a) -> (a -> a -> a) -> a -> a
iter ShapeR sh
shr sh
(sh, Int)
sh' sh -> IO ()
(sh, Int) -> IO ()
write IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
(>>) (() -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ())
      (ArrayData e, e) -> IO (ArrayData e, e)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ArrayData e
aout, e
forall a. HasCallStack => a
undefined)


scanl'Op
    :: forall sh e.
       (e -> e -> e)
    -> e
    -> Delayed (Array (sh, Int) e)
    -> WithReprs (Array (sh, Int) e, Array sh e)
scanl'Op :: forall sh e.
(e -> e -> e)
-> e
-> Delayed (Array (sh, Int) e)
-> WithReprs (Array (sh, Int) e, Array sh e)
scanl'Op e -> e -> e
f e
z (Delayed (ArrayR shr :: ShapeR sh
shr@(ShapeRsnoc ShapeR sh1
shr') TypeR e
tp) (sh1
sh, Int
n) sh -> e
ain Int -> e
_)
  = ( ArrayR (Array (sh, Int) e) -> TupR ArrayR (Array (sh, Int) e)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ShapeR (sh, Int) -> TypeR e -> ArrayR (Array (sh, Int) e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh
ShapeR (sh, Int)
shr TypeR e
TypeR e
tp) TupR ArrayR (Array (sh, Int) e)
-> TupR ArrayR (Array sh e)
-> TupR ArrayR (Array (sh, Int) e, Array sh e)
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` ArrayR (Array sh e) -> TupR ArrayR (Array sh e)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ShapeR sh -> TypeR e -> ArrayR (Array sh e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh
ShapeR sh1
shr' TypeR e
TypeR e
tp)
    , GArrayDataR UniqueArray e
aout GArrayDataR UniqueArray e
-> (Array (sh, Int) e, Array sh e)
-> (Array (sh, Int) e, Array sh e)
forall a b. a -> b -> b
`seq` GArrayDataR UniqueArray e
asum GArrayDataR UniqueArray e
-> (Array (sh, Int) e, Array sh e)
-> (Array (sh, Int) e, Array sh e)
forall a b. a -> b -> b
`seq` ( (sh, Int) -> GArrayDataR UniqueArray e -> Array (sh, Int) e
forall sh e. sh -> ArrayData e -> Array sh e
Array (sh
sh1
sh, Int
n) GArrayDataR UniqueArray e
aout, sh -> GArrayDataR UniqueArray e -> Array sh e
forall sh e. sh -> ArrayData e -> Array sh e
Array sh
sh1
sh GArrayDataR UniqueArray e
asum )
    )
  where
    ((GArrayDataR UniqueArray e
aout, GArrayDataR UniqueArray e
asum), (e, e)
_) = forall e. IO (MutableArrayData e, e) -> (MutableArrayData e, e)
runArrayData @(e, e) (IO (ArrayData (e, e), (e, e)) -> (ArrayData (e, e), (e, e)))
-> IO (ArrayData (e, e), (e, e)) -> (ArrayData (e, e), (e, e))
forall a b. (a -> b) -> a -> b
$ do
      GArrayDataR UniqueArray e
aout <- TypeR e -> Int -> IO (MutableArrayData e)
forall e.
HasCallStack =>
TupR ScalarType e -> Int -> IO (MutableArrayData e)
newArrayData TypeR e
tp (ShapeR sh -> sh -> Int
forall sh. ShapeR sh -> sh -> Int
size ShapeR sh
shr  (sh1
sh, Int
n))
      GArrayDataR UniqueArray e
asum <- TypeR e -> Int -> IO (MutableArrayData e)
forall e.
HasCallStack =>
TupR ScalarType e -> Int -> IO (MutableArrayData e)
newArrayData TypeR e
tp (ShapeR sh1 -> sh1 -> Int
forall sh. ShapeR sh -> sh -> Int
size ShapeR sh1
shr' sh1
sh)

      let write :: (sh1, Int) -> IO ()
write (sh1
sz, Int
0)
            | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0    = TypeR e -> MutableArrayData e -> Int -> e -> IO ()
forall e.
TupR ScalarType e -> MutableArrayData e -> Int -> e -> IO ()
writeArrayData TypeR e
tp GArrayDataR UniqueArray e
MutableArrayData e
asum (ShapeR sh1 -> sh1 -> sh1 -> Int
forall sh. HasCallStack => ShapeR sh -> sh -> sh -> Int
toIndex ShapeR sh1
shr' sh1
sh sh1
sz) e
e
z
            | Bool
otherwise = TypeR e -> MutableArrayData e -> Int -> e -> IO ()
forall e.
TupR ScalarType e -> MutableArrayData e -> Int -> e -> IO ()
writeArrayData TypeR e
tp GArrayDataR UniqueArray e
MutableArrayData e
aout (ShapeR sh -> sh -> sh -> Int
forall sh. HasCallStack => ShapeR sh -> sh -> sh -> Int
toIndex ShapeR sh
shr  (sh1
sh, Int
n) (sh1
sz, Int
0)) e
e
z
          write (sh1
sz, Int
i) = do
            e
x <- TypeR e -> MutableArrayData e -> Int -> IO e
forall e. TupR ScalarType e -> MutableArrayData e -> Int -> IO e
readArrayData TypeR e
tp GArrayDataR UniqueArray e
MutableArrayData e
aout (ShapeR sh -> sh -> sh -> Int
forall sh. HasCallStack => ShapeR sh -> sh -> sh -> Int
toIndex ShapeR sh
shr (sh1
sh, Int
n) (sh1
sz, Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1))
            let y :: e
y = sh -> e
ain (sh1
sz, Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
            if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n
              then TypeR e -> MutableArrayData e -> Int -> e -> IO ()
forall e.
TupR ScalarType e -> MutableArrayData e -> Int -> e -> IO ()
writeArrayData TypeR e
tp GArrayDataR UniqueArray e
MutableArrayData e
asum (ShapeR sh1 -> sh1 -> sh1 -> Int
forall sh. HasCallStack => ShapeR sh -> sh -> sh -> Int
toIndex ShapeR sh1
shr' sh1
sh      sh1
sz)      (e -> e -> e
f e
e
x e
e
y)
              else TypeR e -> MutableArrayData e -> Int -> e -> IO ()
forall e.
TupR ScalarType e -> MutableArrayData e -> Int -> e -> IO ()
writeArrayData TypeR e
tp GArrayDataR UniqueArray e
MutableArrayData e
aout (ShapeR sh -> sh -> sh -> Int
forall sh. HasCallStack => ShapeR sh -> sh -> sh -> Int
toIndex ShapeR sh
shr  (sh1
sh, Int
n) (sh1
sz, Int
i)) (e -> e -> e
f e
e
x e
e
y)

      ShapeR sh
-> sh
-> (sh -> IO ())
-> (IO () -> IO () -> IO ())
-> IO ()
-> IO ()
forall sh a.
ShapeR sh -> sh -> (sh -> a) -> (a -> a -> a) -> a -> a
iter ShapeR sh
shr (sh1
sh, Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) sh -> IO ()
(sh1, Int) -> IO ()
write IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
(>>) (() -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ())
      ((GArrayDataR UniqueArray e, GArrayDataR UniqueArray e), (e, e))
-> IO
     ((GArrayDataR UniqueArray e, GArrayDataR UniqueArray e), (e, e))
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ((GArrayDataR UniqueArray e
aout, GArrayDataR UniqueArray e
asum), (e, e)
forall a. HasCallStack => a
undefined)


scanrOp
    :: forall sh e.
       (e -> e -> e)
    -> e
    -> Delayed (Array (sh, Int) e)
    -> WithReprs (Array (sh, Int) e)
scanrOp :: forall sh e.
(e -> e -> e)
-> e
-> Delayed (Array (sh, Int) e)
-> WithReprs (Array (sh, Int) e)
scanrOp e -> e -> e
f e
z (Delayed (ArrayR ShapeR sh
shr TypeR e
tp) (sh
sz, Int
n) sh -> e
ain Int -> e
_)
  = ( ArrayR (Array (sh, Int) e) -> TupR ArrayR (Array (sh, Int) e)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ShapeR (sh, Int) -> TypeR e -> ArrayR (Array (sh, Int) e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh
ShapeR (sh, Int)
shr TypeR e
TypeR e
tp)
    , ArrayData e
adata ArrayData e -> Array (sh, Int) e -> Array (sh, Int) e
forall a b. a -> b -> b
`seq` (sh, Int) -> ArrayData e -> Array (sh, Int) e
forall sh e. sh -> ArrayData e -> Array sh e
Array (sh, Int)
sh' ArrayData e
adata
    )
  where
    sh' :: (sh, Int)
sh'         = (sh
sz, Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
    --
    (ArrayData e
adata, e
_)  = forall e. IO (MutableArrayData e, e) -> (MutableArrayData e, e)
runArrayData @e (IO (ArrayData e, e) -> (ArrayData e, e))
-> IO (ArrayData e, e) -> (ArrayData e, e)
forall a b. (a -> b) -> a -> b
$ do
      ArrayData e
aout <- TypeR e -> Int -> IO (MutableArrayData e)
forall e.
HasCallStack =>
TupR ScalarType e -> Int -> IO (MutableArrayData e)
newArrayData TypeR e
tp (ShapeR sh -> sh -> Int
forall sh. ShapeR sh -> sh -> Int
size ShapeR sh
shr sh
(sh, Int)
sh')

      let write :: (sh, Int) -> IO ()
write (sh
sz, Int
0) = TypeR e -> MutableArrayData e -> Int -> e -> IO ()
forall e.
TupR ScalarType e -> MutableArrayData e -> Int -> e -> IO ()
writeArrayData TypeR e
tp ArrayData e
MutableArrayData e
aout (ShapeR sh -> sh -> sh -> Int
forall sh. HasCallStack => ShapeR sh -> sh -> sh -> Int
toIndex ShapeR sh
shr sh
(sh, Int)
sh' (sh
sz, Int
n)) e
e
z
          write (sh
sz, Int
i) = do
            let x :: e
x = sh -> e
ain (sh
sz, Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
i)
            e
y <- TypeR e -> MutableArrayData e -> Int -> IO e
forall e. TupR ScalarType e -> MutableArrayData e -> Int -> IO e
readArrayData TypeR e
tp ArrayData e
MutableArrayData e
aout (ShapeR sh -> sh -> sh -> Int
forall sh. HasCallStack => ShapeR sh -> sh -> sh -> Int
toIndex ShapeR sh
shr sh
(sh, Int)
sh' (sh
sz, Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1))
            TypeR e -> MutableArrayData e -> Int -> e -> IO ()
forall e.
TupR ScalarType e -> MutableArrayData e -> Int -> e -> IO ()
writeArrayData TypeR e
tp ArrayData e
MutableArrayData e
aout (ShapeR sh -> sh -> sh -> Int
forall sh. HasCallStack => ShapeR sh -> sh -> sh -> Int
toIndex ShapeR sh
shr sh
(sh, Int)
sh' (sh
sz, Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
i)) (e -> e -> e
f e
e
x e
e
y)

      ShapeR sh
-> sh
-> (sh -> IO ())
-> (IO () -> IO () -> IO ())
-> IO ()
-> IO ()
forall sh a.
ShapeR sh -> sh -> (sh -> a) -> (a -> a -> a) -> a -> a
iter ShapeR sh
shr sh
(sh, Int)
sh' sh -> IO ()
(sh, Int) -> IO ()
write IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
(>>) (() -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ())
      (ArrayData e, e) -> IO (ArrayData e, e)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ArrayData e
aout, e
forall a. HasCallStack => a
undefined)


scanr1Op
    :: forall sh e. HasCallStack
    => (e -> e -> e)
    -> Delayed (Array (sh, Int) e)
    -> WithReprs (Array (sh, Int) e)
scanr1Op :: forall sh e.
HasCallStack =>
(e -> e -> e)
-> Delayed (Array (sh, Int) e) -> WithReprs (Array (sh, Int) e)
scanr1Op e -> e -> e
f (Delayed (ArrayR ShapeR sh
shr TypeR e
tp) sh :: sh
sh@(sh
_, Int
n) sh -> e
ain Int -> e
_)
  = ( ArrayR (Array (sh, Int) e) -> TupR ArrayR (Array (sh, Int) e)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ArrayR (Array (sh, Int) e) -> TupR ArrayR (Array (sh, Int) e))
-> ArrayR (Array (sh, Int) e) -> TupR ArrayR (Array (sh, Int) e)
forall a b. (a -> b) -> a -> b
$ ShapeR (sh, Int) -> TypeR e -> ArrayR (Array (sh, Int) e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh
ShapeR (sh, Int)
shr TypeR e
TypeR e
tp
    , ArrayData e
adata ArrayData e -> Array (sh, Int) e -> Array (sh, Int) e
forall a b. a -> b -> b
`seq` (sh, Int) -> ArrayData e -> Array (sh, Int) e
forall sh e. sh -> ArrayData e -> Array sh e
Array sh
(sh, Int)
sh ArrayData e
adata
    )
  where
    (ArrayData e
adata, e
_)  = forall e. IO (MutableArrayData e, e) -> (MutableArrayData e, e)
runArrayData @e (IO (ArrayData e, e) -> (ArrayData e, e))
-> IO (ArrayData e, e) -> (ArrayData e, e)
forall a b. (a -> b) -> a -> b
$ do
      ArrayData e
aout <- TypeR e -> Int -> IO (MutableArrayData e)
forall e.
HasCallStack =>
TupR ScalarType e -> Int -> IO (MutableArrayData e)
newArrayData TypeR e
tp (ShapeR sh -> sh -> Int
forall sh. ShapeR sh -> sh -> Int
size ShapeR sh
shr sh
sh
sh)

      let write :: (sh, Int) -> IO ()
write (sh
sz, Int
0) = TypeR e -> MutableArrayData e -> Int -> e -> IO ()
forall e.
TupR ScalarType e -> MutableArrayData e -> Int -> e -> IO ()
writeArrayData TypeR e
tp ArrayData e
MutableArrayData e
aout (ShapeR sh -> sh -> sh -> Int
forall sh. HasCallStack => ShapeR sh -> sh -> sh -> Int
toIndex ShapeR sh
shr sh
sh
sh (sh
sz, Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)) (sh -> e
ain (sh
sz, Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1))
          write (sh
sz, Int
i) = do
            let x :: e
x = sh -> e
ain (sh
sz, Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
            e
y <- TypeR e -> MutableArrayData e -> Int -> IO e
forall e. TupR ScalarType e -> MutableArrayData e -> Int -> IO e
readArrayData TypeR e
tp ArrayData e
MutableArrayData e
aout (ShapeR sh -> sh -> sh -> Int
forall sh. HasCallStack => ShapeR sh -> sh -> sh -> Int
toIndex ShapeR sh
shr sh
sh
sh (sh
sz, Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
i))
            TypeR e -> MutableArrayData e -> Int -> e -> IO ()
forall e.
TupR ScalarType e -> MutableArrayData e -> Int -> e -> IO ()
writeArrayData TypeR e
tp ArrayData e
MutableArrayData e
aout (ShapeR sh -> sh -> sh -> Int
forall sh. HasCallStack => ShapeR sh -> sh -> sh -> Int
toIndex ShapeR sh
shr sh
sh
sh (sh
sz, Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)) (e -> e -> e
f e
e
x e
e
y)

      ShapeR sh
-> sh
-> (sh -> IO ())
-> (IO () -> IO () -> IO ())
-> IO ()
-> IO ()
forall sh a.
ShapeR sh -> sh -> (sh -> a) -> (a -> a -> a) -> a -> a
iter ShapeR sh
shr sh
sh
sh sh -> IO ()
(sh, Int) -> IO ()
write IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
(>>) (() -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ())
      (ArrayData e, e) -> IO (ArrayData e, e)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ArrayData e
aout, e
forall a. HasCallStack => a
undefined)


scanr'Op
    :: forall sh e.
       (e -> e -> e)
    -> e
    -> Delayed (Array (sh, Int) e)
    -> WithReprs (Array (sh, Int) e, Array sh e)
scanr'Op :: forall sh e.
(e -> e -> e)
-> e
-> Delayed (Array (sh, Int) e)
-> WithReprs (Array (sh, Int) e, Array sh e)
scanr'Op e -> e -> e
f e
z (Delayed (ArrayR shr :: ShapeR sh
shr@(ShapeRsnoc ShapeR sh1
shr') TypeR e
tp) (sh1
sh, Int
n) sh -> e
ain Int -> e
_)
  = ( ArrayR (Array (sh, Int) e) -> TupR ArrayR (Array (sh, Int) e)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ShapeR (sh, Int) -> TypeR e -> ArrayR (Array (sh, Int) e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh
ShapeR (sh, Int)
shr TypeR e
TypeR e
tp) TupR ArrayR (Array (sh, Int) e)
-> TupR ArrayR (Array sh e)
-> TupR ArrayR (Array (sh, Int) e, Array sh e)
forall (s :: * -> *) a1 b. TupR s a1 -> TupR s b -> TupR s (a1, b)
`TupRpair` ArrayR (Array sh e) -> TupR ArrayR (Array sh e)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ShapeR sh -> TypeR e -> ArrayR (Array sh e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh
ShapeR sh1
shr' TypeR e
TypeR e
tp)
    , GArrayDataR UniqueArray e
aout GArrayDataR UniqueArray e
-> (Array (sh, Int) e, Array sh e)
-> (Array (sh, Int) e, Array sh e)
forall a b. a -> b -> b
`seq` GArrayDataR UniqueArray e
asum GArrayDataR UniqueArray e
-> (Array (sh, Int) e, Array sh e)
-> (Array (sh, Int) e, Array sh e)
forall a b. a -> b -> b
`seq` ( (sh, Int) -> GArrayDataR UniqueArray e -> Array (sh, Int) e
forall sh e. sh -> ArrayData e -> Array sh e
Array (sh
sh1
sh, Int
n) GArrayDataR UniqueArray e
aout, sh -> GArrayDataR UniqueArray e -> Array sh e
forall sh e. sh -> ArrayData e -> Array sh e
Array sh
sh1
sh GArrayDataR UniqueArray e
asum )
    )
  where
    ((GArrayDataR UniqueArray e
aout, GArrayDataR UniqueArray e
asum), (e, e)
_) = forall e. IO (MutableArrayData e, e) -> (MutableArrayData e, e)
runArrayData @(e, e) (IO (ArrayData (e, e), (e, e)) -> (ArrayData (e, e), (e, e)))
-> IO (ArrayData (e, e), (e, e)) -> (ArrayData (e, e), (e, e))
forall a b. (a -> b) -> a -> b
$ do
      GArrayDataR UniqueArray e
aout <- TypeR e -> Int -> IO (MutableArrayData e)
forall e.
HasCallStack =>
TupR ScalarType e -> Int -> IO (MutableArrayData e)
newArrayData TypeR e
tp (ShapeR sh -> sh -> Int
forall sh. ShapeR sh -> sh -> Int
size ShapeR sh
shr  (sh1
sh, Int
n))
      GArrayDataR UniqueArray e
asum <- TypeR e -> Int -> IO (MutableArrayData e)
forall e.
HasCallStack =>
TupR ScalarType e -> Int -> IO (MutableArrayData e)
newArrayData TypeR e
tp (ShapeR sh1 -> sh1 -> Int
forall sh. ShapeR sh -> sh -> Int
size ShapeR sh1
shr' sh1
sh)

      let write :: (sh1, Int) -> IO ()
write (sh1
sz, Int
0)
            | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0    = TypeR e -> MutableArrayData e -> Int -> e -> IO ()
forall e.
TupR ScalarType e -> MutableArrayData e -> Int -> e -> IO ()
writeArrayData TypeR e
tp GArrayDataR UniqueArray e
MutableArrayData e
asum (ShapeR sh1 -> sh1 -> sh1 -> Int
forall sh. HasCallStack => ShapeR sh -> sh -> sh -> Int
toIndex ShapeR sh1
shr' sh1
sh sh1
sz) e
e
z
            | Bool
otherwise = TypeR e -> MutableArrayData e -> Int -> e -> IO ()
forall e.
TupR ScalarType e -> MutableArrayData e -> Int -> e -> IO ()
writeArrayData TypeR e
tp GArrayDataR UniqueArray e
MutableArrayData e
aout (ShapeR sh -> sh -> sh -> Int
forall sh. HasCallStack => ShapeR sh -> sh -> sh -> Int
toIndex ShapeR sh
shr  (sh1
sh, Int
n) (sh1
sz, Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)) e
e
z

          write (sh1
sz, Int
i) = do
            let x :: e
x = sh -> e
ain (sh1
sz, Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
i)
            e
y <- TypeR e -> MutableArrayData e -> Int -> IO e
forall e. TupR ScalarType e -> MutableArrayData e -> Int -> IO e
readArrayData TypeR e
tp GArrayDataR UniqueArray e
MutableArrayData e
aout (ShapeR sh -> sh -> sh -> Int
forall sh. HasCallStack => ShapeR sh -> sh -> sh -> Int
toIndex ShapeR sh
shr (sh1
sh, Int
n) (sh1
sz, Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
i))
            if Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n
              then TypeR e -> MutableArrayData e -> Int -> e -> IO ()
forall e.
TupR ScalarType e -> MutableArrayData e -> Int -> e -> IO ()
writeArrayData TypeR e
tp GArrayDataR UniqueArray e
MutableArrayData e
asum (ShapeR sh1 -> sh1 -> sh1 -> Int
forall sh. HasCallStack => ShapeR sh -> sh -> sh -> Int
toIndex ShapeR sh1
shr' sh1
sh      sh1
sz)          (e -> e -> e
f e
e
x e
e
y)
              else TypeR e -> MutableArrayData e -> Int -> e -> IO ()
forall e.
TupR ScalarType e -> MutableArrayData e -> Int -> e -> IO ()
writeArrayData TypeR e
tp GArrayDataR UniqueArray e
MutableArrayData e
aout (ShapeR sh -> sh -> sh -> Int
forall sh. HasCallStack => ShapeR sh -> sh -> sh -> Int
toIndex ShapeR sh
shr  (sh1
sh, Int
n) (sh1
sz, Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)) (e -> e -> e
f e
e
x e
e
y)

      ShapeR sh
-> sh
-> (sh -> IO ())
-> (IO () -> IO () -> IO ())
-> IO ()
-> IO ()
forall sh a.
ShapeR sh -> sh -> (sh -> a) -> (a -> a -> a) -> a -> a
iter ShapeR sh
shr (sh1
sh, Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) sh -> IO ()
(sh1, Int) -> IO ()
write IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
(>>) (() -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ())
      ((GArrayDataR UniqueArray e, GArrayDataR UniqueArray e), (e, e))
-> IO
     ((GArrayDataR UniqueArray e, GArrayDataR UniqueArray e), (e, e))
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ((GArrayDataR UniqueArray e
aout, GArrayDataR UniqueArray e
asum), (e, e)
forall a. HasCallStack => a
undefined)


permuteOp
    :: forall sh sh' e. HasCallStack
    => (e -> e -> e)
    -> WithReprs (Array sh' e)
    -> (sh -> PrimMaybe sh')
    -> Delayed   (Array sh  e)
    -> WithReprs (Array sh' e)
permuteOp :: forall sh sh' e.
HasCallStack =>
(e -> e -> e)
-> WithReprs (Array sh' e)
-> (sh -> PrimMaybe sh')
-> Delayed (Array sh e)
-> WithReprs (Array sh' e)
permuteOp e -> e -> e
f (TupRsingle (ArrayR ShapeR sh
shr' TypeR e
_), def :: Array sh' e
def@(Array sh'
_ ArrayData e
adef)) sh -> PrimMaybe sh'
p (Delayed (ArrayR ShapeR sh
shr TypeR e
tp) sh
sh sh -> e
_ Int -> e
ain)
  = (ArrayR (Array sh' e) -> ArraysR (Array sh' e)
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle (ArrayR (Array sh' e) -> ArraysR (Array sh' e))
-> ArrayR (Array sh' e) -> ArraysR (Array sh' e)
forall a b. (a -> b) -> a -> b
$ ShapeR sh' -> TypeR e -> ArrayR (Array sh' e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh'
ShapeR sh
shr' TypeR e
TypeR e
tp, ArrayData e
adata ArrayData e -> Array sh' e -> Array sh' e
forall a b. a -> b -> b
`seq` sh' -> ArrayData e -> Array sh' e
forall sh e. sh -> ArrayData e -> Array sh e
Array sh'
sh' ArrayData e
adata)
  where
    sh' :: sh'
sh'         = Array sh' e -> sh'
forall sh e. Array sh e -> sh
shape Array sh' e
def
    n' :: Int
n'          = ShapeR sh -> sh -> Int
forall sh. ShapeR sh -> sh -> Int
size ShapeR sh
shr' sh'
sh
sh'
    --
    (ArrayData e
adata, e
_)  = forall e. IO (MutableArrayData e, e) -> (MutableArrayData e, e)
runArrayData @e (IO (ArrayData e, e) -> (ArrayData e, e))
-> IO (ArrayData e, e) -> (ArrayData e, e)
forall a b. (a -> b) -> a -> b
$ do
      ArrayData e
aout <- TypeR e -> Int -> IO (MutableArrayData e)
forall e.
HasCallStack =>
TupR ScalarType e -> Int -> IO (MutableArrayData e)
newArrayData TypeR e
tp Int
n'

      let -- initialise array with default values
          init :: Int -> IO ()
init Int
i
            | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n'   = () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
            | Bool
otherwise = do
                e
x <- TypeR e -> MutableArrayData e -> Int -> IO e
forall e. TupR ScalarType e -> MutableArrayData e -> Int -> IO e
readArrayData TypeR e
tp ArrayData e
MutableArrayData e
adef Int
i
                TypeR e -> MutableArrayData e -> Int -> e -> IO ()
forall e.
TupR ScalarType e -> MutableArrayData e -> Int -> e -> IO ()
writeArrayData TypeR e
tp ArrayData e
MutableArrayData e
aout Int
i e
x
                Int -> IO ()
init (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)

          -- project each element onto the destination array and update
          update :: sh -> IO ()
update sh
src
            = case sh -> PrimMaybe sh'
p sh
src of
                (PrimBool
0,((), sh')
_)        -> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()
                (PrimBool
1,((),sh'
dst)) -> do
                  let i :: Int
i = ShapeR sh -> sh -> sh -> Int
forall sh. HasCallStack => ShapeR sh -> sh -> sh -> Int
toIndex ShapeR sh
shr  sh
sh
sh  sh
sh
src
                      j :: Int
j = ShapeR sh -> sh -> sh -> Int
forall sh. HasCallStack => ShapeR sh -> sh -> sh -> Int
toIndex ShapeR sh
shr' sh'
sh
sh' sh'
sh
dst
                      x :: e
x = Int -> e
ain Int
i
                  --
                  e
y <- TypeR e -> MutableArrayData e -> Int -> IO e
forall e. TupR ScalarType e -> MutableArrayData e -> Int -> IO e
readArrayData TypeR e
tp ArrayData e
MutableArrayData e
aout Int
j
                  TypeR e -> MutableArrayData e -> Int -> e -> IO ()
forall e.
TupR ScalarType e -> MutableArrayData e -> Int -> e -> IO ()
writeArrayData TypeR e
tp ArrayData e
MutableArrayData e
aout Int
j (e -> e -> e
f e
e
x e
e
y)
                PrimMaybe sh'
_            -> Format (IO ()) (IO ()) -> IO ()
forall r a. HasCallStack => Format r a -> a
internalError Format (IO ()) (IO ())
"unexpected tag"

      Int -> IO ()
init Int
0
      ShapeR sh
-> sh
-> (sh -> IO ())
-> (IO () -> IO () -> IO ())
-> IO ()
-> IO ()
forall sh a.
ShapeR sh -> sh -> (sh -> a) -> (a -> a -> a) -> a -> a
iter ShapeR sh
shr sh
sh
sh sh -> IO ()
sh -> IO ()
update IO () -> IO () -> IO ()
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
(>>) (() -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ())
      (ArrayData e, e) -> IO (ArrayData e, e)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (ArrayData e
aout, e
forall a. HasCallStack => a
undefined)


backpermuteOp
    :: ShapeR sh'
    -> sh'
    -> (sh' -> sh)
    -> Delayed (Array sh e)
    -> WithReprs (Array sh' e)
backpermuteOp :: forall sh' sh e.
ShapeR sh'
-> sh'
-> (sh' -> sh)
-> Delayed (Array sh e)
-> WithReprs (Array sh' e)
backpermuteOp ShapeR sh'
shr sh'
sh' sh' -> sh
p (Delayed (ArrayR ShapeR sh
_ TypeR e
tp) sh
_ sh -> e
arr Int -> e
_)
  = ArrayR (Array sh' e)
-> sh' -> (sh' -> e) -> WithReprs (Array sh' e)
forall sh e.
ArrayR (Array sh e) -> sh -> (sh -> e) -> WithReprs (Array sh e)
fromFunction' (ShapeR sh' -> TypeR e -> ArrayR (Array sh' e)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh'
shr TypeR e
TypeR e
tp) sh'
sh' (\sh'
ix -> sh -> e
arr (sh -> e) -> sh -> e
forall a b. (a -> b) -> a -> b
$ sh' -> sh
p sh'
ix)


stencilOp
    :: HasCallStack
    => StencilR sh a stencil
    -> TypeR b
    -> (stencil -> b)
    -> Boundary (Array sh a)
    -> Delayed  (Array sh a)
    -> WithReprs (Array sh b)
stencilOp :: forall sh a stencil b.
HasCallStack =>
StencilR sh a stencil
-> TypeR b
-> (stencil -> b)
-> Boundary (Array sh a)
-> Delayed (Array sh a)
-> WithReprs (Array sh b)
stencilOp StencilR sh a stencil
stencil TypeR b
tp stencil -> b
f Boundary (Array sh a)
bnd arr :: Delayed (Array sh a)
arr@(Delayed ArrayR (Array sh e)
_ sh
sh sh -> e
_ Int -> e
_)
  = ArrayR (Array sh b) -> sh -> (sh -> b) -> WithReprs (Array sh b)
forall sh e.
ArrayR (Array sh e) -> sh -> (sh -> e) -> WithReprs (Array sh e)
fromFunction' (ShapeR sh -> TypeR b -> ArrayR (Array sh b)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh
shr TypeR b
tp) sh
sh
sh
  ((sh -> b) -> WithReprs (Array sh b))
-> (sh -> b) -> WithReprs (Array sh b)
forall a b. (a -> b) -> a -> b
$ stencil -> b
f (stencil -> b) -> (sh -> stencil) -> sh -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StencilR sh a stencil -> (sh -> a) -> sh -> stencil
forall sh e stencil.
StencilR sh e stencil -> (sh -> e) -> sh -> stencil
stencilAccess StencilR sh a stencil
stencil (ShapeR sh
-> Boundary (Array sh a) -> Delayed (Array sh a) -> sh -> a
forall sh e.
HasCallStack =>
ShapeR sh
-> Boundary (Array sh e) -> Delayed (Array sh e) -> sh -> e
bounded ShapeR sh
shr Boundary (Array sh a)
bnd Delayed (Array sh a)
arr)
  where
    shr :: ShapeR sh
shr = StencilR sh a stencil -> ShapeR sh
forall sh e pat. StencilR sh e pat -> ShapeR sh
stencilShapeR StencilR sh a stencil
stencil


stencil2Op
    :: HasCallStack
    => StencilR sh a stencil1
    -> StencilR sh b stencil2
    -> TypeR c
    -> (stencil1 -> stencil2 -> c)
    -> Boundary (Array sh a)
    -> Delayed  (Array sh a)
    -> Boundary (Array sh b)
    -> Delayed  (Array sh b)
    -> WithReprs (Array sh c)
stencil2Op :: forall sh a stencil1 b stencil2 c.
HasCallStack =>
StencilR sh a stencil1
-> StencilR sh b stencil2
-> TypeR c
-> (stencil1 -> stencil2 -> c)
-> Boundary (Array sh a)
-> Delayed (Array sh a)
-> Boundary (Array sh b)
-> Delayed (Array sh b)
-> WithReprs (Array sh c)
stencil2Op StencilR sh a stencil1
s1 StencilR sh b stencil2
s2 TypeR c
tp stencil1 -> stencil2 -> c
stencil Boundary (Array sh a)
bnd1 arr1 :: Delayed (Array sh a)
arr1@(Delayed ArrayR (Array sh e)
_ sh
sh1 sh -> e
_ Int -> e
_) Boundary (Array sh b)
bnd2 arr2 :: Delayed (Array sh b)
arr2@(Delayed ArrayR (Array sh e)
_ sh
sh2 sh -> e
_ Int -> e
_)
  = ArrayR (Array sh c) -> sh -> (sh -> c) -> WithReprs (Array sh c)
forall sh e.
ArrayR (Array sh e) -> sh -> (sh -> e) -> WithReprs (Array sh e)
fromFunction' (ShapeR sh -> TypeR c -> ArrayR (Array sh c)
forall sh e. ShapeR sh -> TypeR e -> ArrayR (Array sh e)
ArrayR ShapeR sh
shr TypeR c
tp) (ShapeR sh -> sh -> sh -> sh
forall sh. ShapeR sh -> sh -> sh -> sh
intersect ShapeR sh
shr sh
sh
sh1 sh
sh
sh2) sh -> c
f
  where
    f :: sh -> c
f sh
ix  = stencil1 -> stencil2 -> c
stencil (StencilR sh a stencil1 -> (sh -> a) -> sh -> stencil1
forall sh e stencil.
StencilR sh e stencil -> (sh -> e) -> sh -> stencil
stencilAccess StencilR sh a stencil1
s1 (ShapeR sh
-> Boundary (Array sh a) -> Delayed (Array sh a) -> sh -> a
forall sh e.
HasCallStack =>
ShapeR sh
-> Boundary (Array sh e) -> Delayed (Array sh e) -> sh -> e
bounded ShapeR sh
shr Boundary (Array sh a)
bnd1 Delayed (Array sh a)
arr1) sh
ix)
                    (StencilR sh b stencil2 -> (sh -> b) -> sh -> stencil2
forall sh e stencil.
StencilR sh e stencil -> (sh -> e) -> sh -> stencil
stencilAccess StencilR sh b stencil2
s2 (ShapeR sh
-> Boundary (Array sh b) -> Delayed (Array sh b) -> sh -> b
forall sh e.
HasCallStack =>
ShapeR sh
-> Boundary (Array sh e) -> Delayed (Array sh e) -> sh -> e
bounded ShapeR sh
shr Boundary (Array sh b)
bnd2 Delayed (Array sh b)
arr2) sh
ix)
    shr :: ShapeR sh
shr = StencilR sh a stencil1 -> ShapeR sh
forall sh e pat. StencilR sh e pat -> ShapeR sh
stencilShapeR StencilR sh a stencil1
s1

stencilAccess
    :: StencilR sh e stencil
    -> (sh -> e)
    -> sh
    -> stencil
stencilAccess :: forall sh e stencil.
StencilR sh e stencil -> (sh -> e) -> sh -> stencil
stencilAccess StencilR sh e stencil
stencil = ShapeR sh -> StencilR sh e stencil -> (sh -> e) -> sh -> stencil
forall sh e stencil.
ShapeR sh -> StencilR sh e stencil -> (sh -> e) -> sh -> stencil
goR (StencilR sh e stencil -> ShapeR sh
forall sh e pat. StencilR sh e pat -> ShapeR sh
stencilShapeR StencilR sh e stencil
stencil) StencilR sh e stencil
stencil
  where
    -- Base cases, nothing interesting to do here since we know the lower
    -- dimension is Z.
    --
    goR :: ShapeR sh -> StencilR sh e stencil -> (sh -> e) -> sh -> stencil
    goR :: forall sh e stencil.
ShapeR sh -> StencilR sh e stencil -> (sh -> e) -> sh -> stencil
goR ShapeR sh
_ (StencilRunit3 TypeR e
_) sh -> e
rf sh
ix =
      let
          (()
z, Int
i) = sh
ix
          rf' :: Int -> e
rf' Int
d  = sh -> e
rf (()
z, Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
d)
      in
      ((( ()
      , Int -> e
rf' (-Int
1))
      , Int -> e
rf'   Int
0 )
      , Int -> e
rf'   Int
1 )

    goR ShapeR sh
_ (StencilRunit5 TypeR e
_) sh -> e
rf sh
ix =
      let (()
z, Int
i) = sh
ix
          rf' :: Int -> e
rf' Int
d  = sh -> e
rf (()
z, Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
d)
      in
      ((((( ()
      , Int -> e
rf' (-Int
2))
      , Int -> e
rf' (-Int
1))
      , Int -> e
rf'   Int
0 )
      , Int -> e
rf'   Int
1 )
      , Int -> e
rf'   Int
2 )

    goR ShapeR sh
_ (StencilRunit7 TypeR e
_) sh -> e
rf sh
ix =
      let (()
z, Int
i) = sh
ix
          rf' :: Int -> e
rf' Int
d  = sh -> e
rf (()
z, Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
d)
      in
      ((((((( ()
      , Int -> e
rf' (-Int
3))
      , Int -> e
rf' (-Int
2))
      , Int -> e
rf' (-Int
1))
      , Int -> e
rf'   Int
0 )
      , Int -> e
rf'   Int
1 )
      , Int -> e
rf'   Int
2 )
      , Int -> e
rf'   Int
3 )

    goR ShapeR sh
_ (StencilRunit9 TypeR e
_) sh -> e
rf sh
ix =
      let (()
z, Int
i) = sh
ix
          rf' :: Int -> e
rf' Int
d  = sh -> e
rf (()
z, Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
d)
      in
      ((((((((( ()
      , Int -> e
rf' (-Int
4))
      , Int -> e
rf' (-Int
3))
      , Int -> e
rf' (-Int
2))
      , Int -> e
rf' (-Int
1))
      , Int -> e
rf'   Int
0 )
      , Int -> e
rf'   Int
1 )
      , Int -> e
rf'   Int
2 )
      , Int -> e
rf'   Int
3 )
      , Int -> e
rf'   Int
4 )

    -- Recursive cases. Note that because the stencil pattern is defined with
    -- cons ordering, whereas shapes (and indices) are defined as a snoc-list,
    -- when we recurse on the stencil structure we must manipulate the
    -- _left-most_ index component.
    --
    goR (ShapeRsnoc ShapeR sh1
shr) (StencilRtup3 StencilR sh1 e pat1
s1 StencilR sh1 e pat2
s2 StencilR sh1 e pat3
s3) sh -> e
rf sh
ix =
      let (Int
i, sh1
ix') = ShapeR sh1 -> (sh1, Int) -> (Int, sh1)
forall sh. ShapeR sh -> (sh, Int) -> (Int, sh)
uncons ShapeR sh1
shr sh
(sh1, Int)
ix
          rf' :: Int -> sh1 -> e
rf' Int
d sh1
ds = sh -> e
rf (ShapeR sh1 -> Int -> sh1 -> (sh1, Int)
forall sh. ShapeR sh -> Int -> sh -> (sh, Int)
cons ShapeR sh1
shr (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
d) sh1
ds)
      in
      ((( ()
      , ShapeR sh1 -> StencilR sh1 e pat1 -> (sh1 -> e) -> sh1 -> pat1
forall sh e stencil.
ShapeR sh -> StencilR sh e stencil -> (sh -> e) -> sh -> stencil
goR ShapeR sh1
shr StencilR sh1 e pat1
StencilR sh1 e pat1
s1 (Int -> sh1 -> e
rf' (-Int
1)) sh1
ix')
      , ShapeR sh1 -> StencilR sh1 e pat2 -> (sh1 -> e) -> sh1 -> pat2
forall sh e stencil.
ShapeR sh -> StencilR sh e stencil -> (sh -> e) -> sh -> stencil
goR ShapeR sh1
shr StencilR sh1 e pat2
StencilR sh1 e pat2
s2 (Int -> sh1 -> e
rf'   Int
0)  sh1
ix')
      , ShapeR sh1 -> StencilR sh1 e pat3 -> (sh1 -> e) -> sh1 -> pat3
forall sh e stencil.
ShapeR sh -> StencilR sh e stencil -> (sh -> e) -> sh -> stencil
goR ShapeR sh1
shr StencilR sh1 e pat3
StencilR sh1 e pat3
s3 (Int -> sh1 -> e
rf'   Int
1)  sh1
ix')

    goR (ShapeRsnoc ShapeR sh1
shr) (StencilRtup5 StencilR sh1 e pat1
s1 StencilR sh1 e pat2
s2 StencilR sh1 e pat3
s3 StencilR sh1 e pat4
s4 StencilR sh1 e pat5
s5) sh -> e
rf sh
ix =
      let (Int
i, sh1
ix') = ShapeR sh1 -> (sh1, Int) -> (Int, sh1)
forall sh. ShapeR sh -> (sh, Int) -> (Int, sh)
uncons ShapeR sh1
shr sh
(sh1, Int)
ix
          rf' :: Int -> sh1 -> e
rf' Int
d sh1
ds = sh -> e
rf (ShapeR sh1 -> Int -> sh1 -> (sh1, Int)
forall sh. ShapeR sh -> Int -> sh -> (sh, Int)
cons ShapeR sh1
shr (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
d) sh1
ds)
      in
      ((((( ()
      , ShapeR sh1 -> StencilR sh1 e pat1 -> (sh1 -> e) -> sh1 -> pat1
forall sh e stencil.
ShapeR sh -> StencilR sh e stencil -> (sh -> e) -> sh -> stencil
goR ShapeR sh1
shr StencilR sh1 e pat1
StencilR sh1 e pat1
s1 (Int -> sh1 -> e
rf' (-Int
2)) sh1
ix')
      , ShapeR sh1 -> StencilR sh1 e pat2 -> (sh1 -> e) -> sh1 -> pat2
forall sh e stencil.
ShapeR sh -> StencilR sh e stencil -> (sh -> e) -> sh -> stencil
goR ShapeR sh1
shr StencilR sh1 e pat2
StencilR sh1 e pat2
s2 (Int -> sh1 -> e
rf' (-Int
1)) sh1
ix')
      , ShapeR sh1 -> StencilR sh1 e pat3 -> (sh1 -> e) -> sh1 -> pat3
forall sh e stencil.
ShapeR sh -> StencilR sh e stencil -> (sh -> e) -> sh -> stencil
goR ShapeR sh1
shr StencilR sh1 e pat3
StencilR sh1 e pat3
s3 (Int -> sh1 -> e
rf'   Int
0)  sh1
ix')
      , ShapeR sh1 -> StencilR sh1 e pat4 -> (sh1 -> e) -> sh1 -> pat4
forall sh e stencil.
ShapeR sh -> StencilR sh e stencil -> (sh -> e) -> sh -> stencil
goR ShapeR sh1
shr StencilR sh1 e pat4
StencilR sh1 e pat4
s4 (Int -> sh1 -> e
rf'   Int
1)  sh1
ix')
      , ShapeR sh1 -> StencilR sh1 e pat5 -> (sh1 -> e) -> sh1 -> pat5
forall sh e stencil.
ShapeR sh -> StencilR sh e stencil -> (sh -> e) -> sh -> stencil
goR ShapeR sh1
shr StencilR sh1 e pat5
StencilR sh1 e pat5
s5 (Int -> sh1 -> e
rf'   Int
2)  sh1
ix')

    goR (ShapeRsnoc ShapeR sh1
shr) (StencilRtup7 StencilR sh1 e pat1
s1 StencilR sh1 e pat2
s2 StencilR sh1 e pat3
s3 StencilR sh1 e pat4
s4 StencilR sh1 e pat5
s5 StencilR sh1 e pat6
s6 StencilR sh1 e pat7
s7) sh -> e
rf sh
ix =
      let (Int
i, sh1
ix') = ShapeR sh1 -> (sh1, Int) -> (Int, sh1)
forall sh. ShapeR sh -> (sh, Int) -> (Int, sh)
uncons ShapeR sh1
shr sh
(sh1, Int)
ix
          rf' :: Int -> sh1 -> e
rf' Int
d sh1
ds = sh -> e
rf (ShapeR sh1 -> Int -> sh1 -> (sh1, Int)
forall sh. ShapeR sh -> Int -> sh -> (sh, Int)
cons ShapeR sh1
shr (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
d) sh1
ds)
      in
      ((((((( ()
      , ShapeR sh1 -> StencilR sh1 e pat1 -> (sh1 -> e) -> sh1 -> pat1
forall sh e stencil.
ShapeR sh -> StencilR sh e stencil -> (sh -> e) -> sh -> stencil
goR ShapeR sh1
shr StencilR sh1 e pat1
StencilR sh1 e pat1
s1 (Int -> sh1 -> e
rf' (-Int
3)) sh1
ix')
      , ShapeR sh1 -> StencilR sh1 e pat2 -> (sh1 -> e) -> sh1 -> pat2
forall sh e stencil.
ShapeR sh -> StencilR sh e stencil -> (sh -> e) -> sh -> stencil
goR ShapeR sh1
shr StencilR sh1 e pat2
StencilR sh1 e pat2
s2 (Int -> sh1 -> e
rf' (-Int
2)) sh1
ix')
      , ShapeR sh1 -> StencilR sh1 e pat3 -> (sh1 -> e) -> sh1 -> pat3
forall sh e stencil.
ShapeR sh -> StencilR sh e stencil -> (sh -> e) -> sh -> stencil
goR ShapeR sh1
shr StencilR sh1 e pat3
StencilR sh1 e pat3
s3 (Int -> sh1 -> e
rf' (-Int
1)) sh1
ix')
      , ShapeR sh1 -> StencilR sh1 e pat4 -> (sh1 -> e) -> sh1 -> pat4
forall sh e stencil.
ShapeR sh -> StencilR sh e stencil -> (sh -> e) -> sh -> stencil
goR ShapeR sh1
shr StencilR sh1 e pat4
StencilR sh1 e pat4
s4 (Int -> sh1 -> e
rf'   Int
0)  sh1
ix')
      , ShapeR sh1 -> StencilR sh1 e pat5 -> (sh1 -> e) -> sh1 -> pat5
forall sh e stencil.
ShapeR sh -> StencilR sh e stencil -> (sh -> e) -> sh -> stencil
goR ShapeR sh1
shr StencilR sh1 e pat5
StencilR sh1 e pat5
s5 (Int -> sh1 -> e
rf'   Int
1)  sh1
ix')
      , ShapeR sh1 -> StencilR sh1 e pat6 -> (sh1 -> e) -> sh1 -> pat6
forall sh e stencil.
ShapeR sh -> StencilR sh e stencil -> (sh -> e) -> sh -> stencil
goR ShapeR sh1
shr StencilR sh1 e pat6
StencilR sh1 e pat6
s6 (Int -> sh1 -> e
rf'   Int
2)  sh1
ix')
      , ShapeR sh1 -> StencilR sh1 e pat7 -> (sh1 -> e) -> sh1 -> pat7
forall sh e stencil.
ShapeR sh -> StencilR sh e stencil -> (sh -> e) -> sh -> stencil
goR ShapeR sh1
shr StencilR sh1 e pat7
StencilR sh1 e pat7
s7 (Int -> sh1 -> e
rf'   Int
3)  sh1
ix')

    goR (ShapeRsnoc ShapeR sh1
shr) (StencilRtup9 StencilR sh1 e pat1
s1 StencilR sh1 e pat2
s2 StencilR sh1 e pat3
s3 StencilR sh1 e pat4
s4 StencilR sh1 e pat5
s5 StencilR sh1 e pat6
s6 StencilR sh1 e pat7
s7 StencilR sh1 e pat8
s8 StencilR sh1 e pat9
s9) sh -> e
rf sh
ix =
      let (Int
i, sh1
ix') = ShapeR sh1 -> (sh1, Int) -> (Int, sh1)
forall sh. ShapeR sh -> (sh, Int) -> (Int, sh)
uncons ShapeR sh1
shr sh
(sh1, Int)
ix
          rf' :: Int -> sh1 -> e
rf' Int
d sh1
ds = sh -> e
rf (ShapeR sh1 -> Int -> sh1 -> (sh1, Int)
forall sh. ShapeR sh -> Int -> sh -> (sh, Int)
cons ShapeR sh1
shr (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
d) sh1
ds)
      in
      ((((((((( ()
      , ShapeR sh1 -> StencilR sh1 e pat1 -> (sh1 -> e) -> sh1 -> pat1
forall sh e stencil.
ShapeR sh -> StencilR sh e stencil -> (sh -> e) -> sh -> stencil
goR ShapeR sh1
shr StencilR sh1 e pat1
StencilR sh1 e pat1
s1 (Int -> sh1 -> e
rf' (-Int
4)) sh1
ix')
      , ShapeR sh1 -> StencilR sh1 e pat2 -> (sh1 -> e) -> sh1 -> pat2
forall sh e stencil.
ShapeR sh -> StencilR sh e stencil -> (sh -> e) -> sh -> stencil
goR ShapeR sh1
shr StencilR sh1 e pat2
StencilR sh1 e pat2
s2 (Int -> sh1 -> e
rf' (-Int
3)) sh1
ix')
      , ShapeR sh1 -> StencilR sh1 e pat3 -> (sh1 -> e) -> sh1 -> pat3
forall sh e stencil.
ShapeR sh -> StencilR sh e stencil -> (sh -> e) -> sh -> stencil
goR ShapeR sh1
shr StencilR sh1 e pat3
StencilR sh1 e pat3
s3 (Int -> sh1 -> e
rf' (-Int
2)) sh1
ix')
      , ShapeR sh1 -> StencilR sh1 e pat4 -> (sh1 -> e) -> sh1 -> pat4
forall sh e stencil.
ShapeR sh -> StencilR sh e stencil -> (sh -> e) -> sh -> stencil
goR ShapeR sh1
shr StencilR sh1 e pat4
StencilR sh1 e pat4
s4 (Int -> sh1 -> e
rf' (-Int
1)) sh1
ix')
      , ShapeR sh1 -> StencilR sh1 e pat5 -> (sh1 -> e) -> sh1 -> pat5
forall sh e stencil.
ShapeR sh -> StencilR sh e stencil -> (sh -> e) -> sh -> stencil
goR ShapeR sh1
shr StencilR sh1 e pat5
StencilR sh1 e pat5
s5 (Int -> sh1 -> e
rf'   Int
0)  sh1
ix')
      , ShapeR sh1 -> StencilR sh1 e pat6 -> (sh1 -> e) -> sh1 -> pat6
forall sh e stencil.
ShapeR sh -> StencilR sh e stencil -> (sh -> e) -> sh -> stencil
goR ShapeR sh1
shr StencilR sh1 e pat6
StencilR sh1 e pat6
s6 (Int -> sh1 -> e
rf'   Int
1)  sh1
ix')
      , ShapeR sh1 -> StencilR sh1 e pat7 -> (sh1 -> e) -> sh1 -> pat7
forall sh e stencil.
ShapeR sh -> StencilR sh e stencil -> (sh -> e) -> sh -> stencil
goR ShapeR sh1
shr StencilR sh1 e pat7
StencilR sh1 e pat7
s7 (Int -> sh1 -> e
rf'   Int
2)  sh1
ix')
      , ShapeR sh1 -> StencilR sh1 e pat8 -> (sh1 -> e) -> sh1 -> pat8
forall sh e stencil.
ShapeR sh -> StencilR sh e stencil -> (sh -> e) -> sh -> stencil
goR ShapeR sh1
shr StencilR sh1 e pat8
StencilR sh1 e pat8
s8 (Int -> sh1 -> e
rf'   Int
3)  sh1
ix')
      , ShapeR sh1 -> StencilR sh1 e pat9 -> (sh1 -> e) -> sh1 -> pat9
forall sh e stencil.
ShapeR sh -> StencilR sh e stencil -> (sh -> e) -> sh -> stencil
goR ShapeR sh1
shr StencilR sh1 e pat9
StencilR sh1 e pat9
s9 (Int -> sh1 -> e
rf'   Int
4)  sh1
ix')

    -- Add a left-most component to an index
    --
    cons :: ShapeR sh -> Int -> sh -> (sh, Int)
    cons :: forall sh. ShapeR sh -> Int -> sh -> (sh, Int)
cons ShapeR sh
ShapeRz          Int
ix ()       = ((), Int
ix)
    cons (ShapeRsnoc ShapeR sh1
shr) Int
ix (sh1
sh, Int
sz) = (ShapeR sh1 -> Int -> sh1 -> (sh1, Int)
forall sh. ShapeR sh -> Int -> sh -> (sh, Int)
cons ShapeR sh1
shr Int
ix sh1
sh, Int
sz)

    -- Remove the left-most index of an index, and return the remainder
    --
    uncons :: ShapeR sh -> (sh, Int) -> (Int, sh)
    uncons :: forall sh. ShapeR sh -> (sh, Int) -> (Int, sh)
uncons ShapeR sh
ShapeRz          ((), Int
v)  = (Int
v, ())
    uncons (ShapeRsnoc ShapeR sh1
shr) (sh
v1, Int
v2) = let (Int
i, sh1
v1') = ShapeR sh1 -> (sh1, Int) -> (Int, sh1)
forall sh. ShapeR sh -> (sh, Int) -> (Int, sh)
uncons ShapeR sh1
shr sh
(sh1, Int)
v1
                                       in  (Int
i, (sh1
v1', Int
v2))


bounded
    :: HasCallStack
    => ShapeR sh
    -> Boundary (Array sh e)
    -> Delayed (Array sh e)
    -> sh
    -> e
bounded :: forall sh e.
HasCallStack =>
ShapeR sh
-> Boundary (Array sh e) -> Delayed (Array sh e) -> sh -> e
bounded ShapeR sh
shr Boundary (Array sh e)
bnd (Delayed ArrayR (Array sh e)
_ sh
sh sh -> e
f Int -> e
_) sh
ix =
  if ShapeR sh -> sh -> sh -> Bool
forall sh. ShapeR sh -> sh -> sh -> Bool
inside ShapeR sh
shr sh
sh
sh sh
ix
    then sh -> e
f sh
sh
ix
    else
      case Boundary (Array sh e)
bnd of
        Function sh -> e
g -> sh -> e
g sh
sh
ix
        Constant t
v -> e
t
v
        Boundary (Array sh e)
_          -> sh -> e
f (ShapeR sh -> sh -> sh -> sh
forall sh. HasCallStack => ShapeR sh -> sh -> sh -> sh
bound ShapeR sh
ShapeR sh
shr sh
sh sh
sh
ix)

  where
    -- Whether the index (second argument) is inside the bounds of the given
    -- shape (first argument).
    --
    inside :: ShapeR sh -> sh -> sh -> Bool
    inside :: forall sh. ShapeR sh -> sh -> sh -> Bool
inside ShapeR sh
ShapeRz          ()       ()       = Bool
True
    inside (ShapeRsnoc ShapeR sh1
shr) (sh1
sh, Int
sz) (sh1
ih, Int
iz) = Int
iz Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
iz Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz Bool -> Bool -> Bool
&& ShapeR sh1 -> sh1 -> sh1 -> Bool
forall sh. ShapeR sh -> sh -> sh -> Bool
inside ShapeR sh1
shr sh1
sh sh1
ih

    -- Return the index (second argument), updated to obey the given boundary
    -- conditions when outside the bounds of the given shape (first argument)
    --
    bound :: HasCallStack => ShapeR sh -> sh -> sh -> sh
    bound :: forall sh. HasCallStack => ShapeR sh -> sh -> sh -> sh
bound ShapeR sh
ShapeRz () () = ()
    bound (ShapeRsnoc ShapeR sh1
shr) (sh1
sh, Int
sz) (sh1
ih, Int
iz) = (ShapeR sh1 -> sh1 -> sh1 -> sh1
forall sh. HasCallStack => ShapeR sh -> sh -> sh -> sh
bound ShapeR sh1
shr sh1
sh sh1
ih, Int
ih')
      where
        ih' :: Int
ih'
          | Int
iz Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = case Boundary (Array sh e)
bnd of
                          Boundary (Array sh e)
Clamp  -> Int
0
                          Boundary (Array sh e)
Mirror -> -Int
iz
                          Boundary (Array sh e)
Wrap   -> Int
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
iz
                          Boundary (Array sh e)
_      -> Format Int Int -> Int
forall r a. HasCallStack => Format r a -> a
internalError Format Int Int
"unexpected boundary condition"
          | Int
iz Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
sz  = case Boundary (Array sh e)
bnd of
                          Boundary (Array sh e)
Clamp  -> Int
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
                          Boundary (Array sh e)
Mirror -> Int
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
- (Int
iz Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2)
                          Boundary (Array sh e)
Wrap   -> Int
iz Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
sz
                          Boundary (Array sh e)
_      -> Format Int Int -> Int
forall r a. HasCallStack => Format r a -> a
internalError Format Int Int
"unexpected boundary condition"
          | Bool
otherwise = Int
iz

-- toSeqOp :: forall slix sl dim co e proxy. (Elt slix, Shape sl, Shape dim, Elt e)
--         => SliceIndex (EltRepr slix)
--                       (EltRepr sl)
--                       co
--                       (EltRepr dim)
--         -> proxy slix
--         -> Array dim e
--         -> [Array sl e]
-- toSeqOp sliceIndex _ arr = map (sliceOp sliceIndex arr :: slix -> Array sl e)
--                                (enumSlices sliceIndex (shape arr))


-- Stencil boundary conditions
-- ---------------------------

data Boundary t where
  Clamp    :: Boundary t
  Mirror   :: Boundary t
  Wrap     :: Boundary t
  Constant :: t -> Boundary (Array sh t)
  Function :: (sh -> e) -> Boundary (Array sh e)


evalBoundary :: HasCallStack => AST.Boundary aenv t -> Val aenv -> Boundary t
evalBoundary :: forall aenv t.
HasCallStack =>
Boundary aenv t -> Val aenv -> Boundary t
evalBoundary Boundary aenv t
bnd Val aenv
aenv =
  case Boundary aenv t
bnd of
    Boundary aenv t
AST.Clamp      -> Boundary t
forall t. Boundary t
Clamp
    Boundary aenv t
AST.Mirror     -> Boundary t
forall t. Boundary t
Mirror
    Boundary aenv t
AST.Wrap       -> Boundary t
forall t. Boundary t
Wrap
    AST.Constant e
v -> e -> Boundary (Array sh e)
forall sh e. sh -> Boundary (Array e sh)
Constant e
v
    AST.Function Fun aenv (sh -> e)
f -> (sh -> e) -> Boundary (Array sh e)
forall sh e. (sh -> e) -> Boundary (Array sh e)
Function (Fun aenv (sh -> e) -> Val aenv -> sh -> e
forall aenv t. HasCallStack => Fun aenv t -> Val aenv -> t
evalFun Fun aenv (sh -> e)
f Val aenv
aenv)

atraceOp :: Message as -> as -> IO ()
atraceOp :: forall as. Message as -> as -> IO ()
atraceOp (Message as -> [Char]
show Maybe (CodeQ (as -> [Char]))
_ Text
msg) as
as =
  let str :: [Char]
str = as -> [Char]
show as
as
   in do
     if [Char] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Char]
str
        then Handle -> Text -> IO ()
T.hPutStrLn Handle
stderr Text
msg
        else Handle
-> Format (IO ()) (Text -> [Char] -> IO ())
-> Text
-> [Char]
-> IO ()
forall (m :: * -> *) a. MonadIO m => Handle -> Format (m ()) a -> a
hprint Handle
stderr (Format ([Char] -> IO ()) (Text -> [Char] -> IO ())
forall r. Format r (Text -> r)
stext Format ([Char] -> IO ()) (Text -> [Char] -> IO ())
-> Format (IO ()) ([Char] -> IO ())
-> Format (IO ()) (Text -> [Char] -> IO ())
forall r a r'. Format r a -> Format r' r -> Format r' a
% Format ([Char] -> IO ()) ([Char] -> IO ())
": " Format ([Char] -> IO ()) ([Char] -> IO ())
-> Format (IO ()) ([Char] -> IO ())
-> Format (IO ()) ([Char] -> IO ())
forall r a r'. Format r a -> Format r' r -> Format r' a
% Format (IO ()) ([Char] -> IO ())
forall r. Format r ([Char] -> r)
string Format (IO ()) ([Char] -> IO ())
-> Format (IO ()) (IO ()) -> Format (IO ()) ([Char] -> IO ())
forall r a r'. Format r a -> Format r' r -> Format r' a
% Format (IO ()) (IO ())
"\n") Text
msg [Char]
str
     Handle -> IO ()
hFlush Handle
stderr


-- Scalar expression evaluation
-- ----------------------------

-- Evaluate a closed scalar expression
--
evalExp :: HasCallStack => Exp aenv t -> Val aenv -> t
evalExp :: forall aenv t. HasCallStack => Exp aenv t -> Val aenv -> t
evalExp Exp aenv t
e Val aenv
aenv = Exp aenv t -> Val () -> Val aenv -> t
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> Val env -> Val aenv -> t
evalOpenExp Exp aenv t
e Val ()
Empty Val aenv
aenv

-- Evaluate a closed scalar function
--
evalFun :: HasCallStack => Fun aenv t -> Val aenv -> t
evalFun :: forall aenv t. HasCallStack => Fun aenv t -> Val aenv -> t
evalFun Fun aenv t
f Val aenv
aenv = Fun aenv t -> Val () -> Val aenv -> t
forall env aenv t.
HasCallStack =>
OpenFun env aenv t -> Val env -> Val aenv -> t
evalOpenFun Fun aenv t
f Val ()
Empty Val aenv
aenv

-- Evaluate an open scalar function
--
evalOpenFun :: HasCallStack => OpenFun env aenv t -> Val env -> Val aenv -> t
evalOpenFun :: forall env aenv t.
HasCallStack =>
OpenFun env aenv t -> Val env -> Val aenv -> t
evalOpenFun (Body OpenExp env aenv t
e)    Val env
env Val aenv
aenv = OpenExp env aenv t -> Val env -> Val aenv -> t
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> Val env -> Val aenv -> t
evalOpenExp OpenExp env aenv t
e Val env
env Val aenv
aenv
evalOpenFun (Lam ELeftHandSide a env env'
lhs OpenFun env' aenv t1
f) Val env
env Val aenv
aenv =
  \a
x -> OpenFun env' aenv t1 -> Val env' -> Val aenv -> t1
forall env aenv t.
HasCallStack =>
OpenFun env aenv t -> Val env -> Val aenv -> t
evalOpenFun OpenFun env' aenv t1
f (Val env
env Val env -> (ELeftHandSide a env env', a) -> Val env'
forall env (s :: * -> *) t env'.
Val env -> (LeftHandSide s t env env', t) -> Val env'
`push` (ELeftHandSide a env env'
lhs, a
x)) Val aenv
aenv


-- Evaluate an open scalar expression
--
-- NB: The implementation of 'Index' and 'Shape' demonstrate clearly why
--     array expressions must be hoisted out of scalar expressions before code
--     execution. If these operations are in the body of a function that gets
--     mapped over an array, the array argument would be evaluated many times
--     leading to a large amount of wasteful recomputation.
--
evalOpenExp
    :: forall env aenv t. HasCallStack
    => OpenExp env aenv t
    -> Val env
    -> Val aenv
    -> t
evalOpenExp :: forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> Val env -> Val aenv -> t
evalOpenExp OpenExp env aenv t
pexp Val env
env Val aenv
aenv =
  let
      evalE :: OpenExp env aenv t' -> t'
      evalE :: forall t'. OpenExp env aenv t' -> t'
evalE OpenExp env aenv t'
e = OpenExp env aenv t' -> Val env -> Val aenv -> t'
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> Val env -> Val aenv -> t
evalOpenExp OpenExp env aenv t'
e Val env
env Val aenv
aenv

      evalF :: OpenFun env aenv f' -> f'
      evalF :: forall f'. OpenFun env aenv f' -> f'
evalF OpenFun env aenv f'
f = OpenFun env aenv f' -> Val env -> Val aenv -> f'
forall env aenv t.
HasCallStack =>
OpenFun env aenv t -> Val env -> Val aenv -> t
evalOpenFun OpenFun env aenv f'
f Val env
env Val aenv
aenv

      evalA :: ArrayVar aenv a -> WithReprs a
      evalA :: forall a. ArrayVar aenv a -> WithReprs a
evalA (Var ArrayR a
repr Idx aenv a
ix) = (ArrayR a -> TupR ArrayR a
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle ArrayR a
repr, Idx aenv a -> Val aenv -> a
forall env t. Idx env t -> Val env -> t
prj Idx aenv a
ix Val aenv
aenv)
  in
  case OpenExp env aenv t
pexp of
    Let ELeftHandSide bnd_t env env'
lhs OpenExp env aenv bnd_t
exp1 OpenExp env' aenv t
exp2           -> let !v1 :: bnd_t
v1  = OpenExp env aenv bnd_t -> bnd_t
forall t'. OpenExp env aenv t' -> t'
evalE OpenExp env aenv bnd_t
exp1
                                       env' :: Val env'
env' = Val env
env Val env -> (ELeftHandSide bnd_t env env', bnd_t) -> Val env'
forall env (s :: * -> *) t env'.
Val env -> (LeftHandSide s t env env', t) -> Val env'
`push` (ELeftHandSide bnd_t env env'
lhs, bnd_t
v1)
                                   in  OpenExp env' aenv t -> Val env' -> Val aenv -> t
forall env aenv t.
HasCallStack =>
OpenExp env aenv t -> Val env -> Val aenv -> t
evalOpenExp OpenExp env' aenv t
exp2 Val env'
env' Val aenv
aenv
    Evar (Var ScalarType t
_ Idx env t
ix)             -> Idx env t -> Val env -> t
forall env t. Idx env t -> Val env -> t
prj Idx env t
ix Val env
env
    Const ScalarType t
_ t
c                   -> t
c
    Undef ScalarType t
tp                    -> TypeR t -> t
forall t. TypeR t -> t
undefElt (ScalarType t -> TypeR t
forall (s :: * -> *) a. s a -> TupR s a
TupRsingle ScalarType t
tp)
    PrimConst PrimConst t
c                 -> PrimConst t -> t
forall a. PrimConst a -> a
evalPrimConst PrimConst t
c
    PrimApp PrimFun (a -> t)
f OpenExp env aenv a
x                 -> PrimFun (a -> t) -> a -> t
forall a r. PrimFun (a -> r) -> a -> r
evalPrim PrimFun (a -> t)
f (OpenExp env aenv a -> a
forall t'. OpenExp env aenv t' -> t'
evalE OpenExp env aenv a
x)
    OpenExp env aenv t
Nil                         -> ()
    Pair OpenExp env aenv t1
e1 OpenExp env aenv t2
e2                  -> let !x1 :: t1
x1 = OpenExp env aenv t1 -> t1
forall t'. OpenExp env aenv t' -> t'
evalE OpenExp env aenv t1
e1
                                       !x2 :: t2
x2 = OpenExp env aenv t2 -> t2
forall t'. OpenExp env aenv t' -> t'
evalE OpenExp env aenv t2
e2
                                   in  (t1
x1, t2
x2)
    VecPack   VecR n s tup
vecR OpenExp env aenv tup
e            -> VecR n s tup -> tup -> Vec n s
forall (n :: Nat) single tuple.
KnownNat n =>
VecR n single tuple -> tuple -> Vec n single
pack   VecR n s tup
vecR (tup -> t) -> tup -> t
forall a b. (a -> b) -> a -> b
$! OpenExp env aenv tup -> tup
forall t'. OpenExp env aenv t' -> t'
evalE OpenExp env aenv tup
e
    VecUnpack VecR n s t
vecR OpenExp env aenv (Vec n s)
e            -> VecR n s t -> Vec n s -> t
forall (n :: Nat) single tuple.
KnownNat n =>
VecR n single tuple -> Vec n single -> tuple
unpack VecR n s t
vecR (Vec n s -> t) -> Vec n s -> t
forall a b. (a -> b) -> a -> b
$! OpenExp env aenv (Vec n s) -> Vec n s
forall t'. OpenExp env aenv t' -> t'
evalE OpenExp env aenv (Vec n s)
e
    IndexSlice SliceIndex slix t co sh
slice OpenExp env aenv slix
slix OpenExp env aenv sh
sh    -> SliceIndex slix t co sh -> slix -> sh -> t
forall slix sl co sh. SliceIndex slix sl co sh -> slix -> sh -> sl
restrict SliceIndex slix t co sh
slice (OpenExp env aenv slix -> slix
forall t'. OpenExp env aenv t' -> t'
evalE OpenExp env aenv slix
slix)
                                                  (OpenExp env aenv sh -> sh
forall t'. OpenExp env aenv t' -> t'
evalE OpenExp env aenv sh
sh)
      where
        restrict :: SliceIndex slix sl co sh -> slix -> sh -> sl
        restrict :: forall slix sl co sh. SliceIndex slix sl co sh -> slix -> sh -> sl
restrict SliceIndex slix sl co sh
SliceNil              ()        ()         = ()
        restrict (SliceAll SliceIndex ix1 slice1 co dim
sliceIdx)   (ix1
slx, ()) (dim
sl, Int
sz)   =
          let sl' :: slice1
sl' = SliceIndex ix1 slice1 co dim -> ix1 -> dim -> slice1
forall slix sl co sh. SliceIndex slix sl co sh -> slix -> sh -> sl
restrict SliceIndex ix1 slice1 co dim
sliceIdx ix1
slx dim
sl
          in  (slice1
sl', Int
sz)
        restrict (SliceFixed SliceIndex ix1 sl co dim
sliceIdx) (ix1
slx, Int
_i)  (dim
sl, Int
_sz) =
          SliceIndex ix1 sl co dim -> ix1 -> dim -> sl
forall slix sl co sh. SliceIndex slix sl co sh -> slix -> sh -> sl
restrict SliceIndex ix1 sl co dim
sliceIdx ix1
slx dim
sl

    IndexFull SliceIndex slix sl co t
slice OpenExp env aenv slix
slix OpenExp env aenv sl
sh     -> SliceIndex slix sl co t -> slix -> sl -> t
forall slix sl co sh. SliceIndex slix sl co sh -> slix -> sl -> sh
extend SliceIndex slix sl co t
slice (OpenExp env aenv slix -> slix
forall t'. OpenExp env aenv t' -> t'
evalE OpenExp env aenv slix
slix)
                                                (OpenExp env aenv sl -> sl
forall t'. OpenExp env aenv t' -> t'
evalE OpenExp env aenv sl
sh)
      where
        extend :: SliceIndex slix sl co sh -> slix -> sl -> sh
        extend :: forall slix sl co sh. SliceIndex slix sl co sh -> slix -> sl -> sh
extend SliceIndex slix sl co sh
SliceNil              ()        ()       = ()
        extend (SliceAll SliceIndex ix1 slice1 co dim
sliceIdx)   (ix1
slx, ()) (slice1
sl, Int
sz) =
          let sh' :: dim
sh' = SliceIndex ix1 slice1 co dim -> ix1 -> slice1 -> dim
forall slix sl co sh. SliceIndex slix sl co sh -> slix -> sl -> sh
extend SliceIndex ix1 slice1 co dim
sliceIdx ix1
slx slice1
sl
          in  (dim
sh', Int
sz)
        extend (SliceFixed SliceIndex ix1 sl co dim
sliceIdx) (ix1
slx, Int
sz) sl
sl       =
          let sh' :: dim
sh' = SliceIndex ix1 sl co dim -> ix1 -> sl -> dim
forall slix sl co sh. SliceIndex slix sl co sh -> slix -> sl -> sh
extend SliceIndex ix1 sl co dim
sliceIdx ix1
slx sl
sl
          in  (dim
sh', Int
sz)

    ToIndex ShapeR sh
shr OpenExp env aenv sh
sh OpenExp env aenv sh
ix           -> ShapeR sh -> sh -> sh -> Int
forall sh. HasCallStack => ShapeR sh -> sh -> sh -> Int
toIndex ShapeR sh
shr (OpenExp env aenv sh -> sh
forall t'. OpenExp env aenv t' -> t'
evalE OpenExp env aenv sh
sh) (OpenExp env aenv sh -> sh
forall t'. OpenExp env aenv t' -> t'
evalE OpenExp env aenv sh
ix)
    FromIndex ShapeR t
shr OpenExp env aenv t
sh OpenExp env aenv Int
ix         -> ShapeR t -> t -> Int -> t
forall sh. HasCallStack => ShapeR sh -> sh -> Int -> sh
fromIndex ShapeR t
shr (OpenExp env aenv t -> t
forall t'. OpenExp env aenv t' -> t'
evalE OpenExp env aenv t
sh) (OpenExp env aenv Int -> Int
forall t'. OpenExp env aenv t' -> t'
evalE OpenExp env aenv Int
ix)
    Case OpenExp env aenv PrimBool
e [(PrimBool, OpenExp env aenv t)]
rhs Maybe (OpenExp env aenv t)
def              -> OpenExp env aenv t -> t
forall t'. OpenExp env aenv t' -> t'
evalE (PrimBool -> [(PrimBool, OpenExp env aenv t)] -> OpenExp env aenv t
caseof (OpenExp env aenv PrimBool -> PrimBool
forall t'. OpenExp env aenv t' -> t'
evalE OpenExp env aenv PrimBool
e) [(PrimBool, OpenExp env aenv t)]
rhs)
      where
        caseof :: TAG -> [(TAG, OpenExp env aenv t)] -> OpenExp env aenv t
        caseof :: PrimBool -> [(PrimBool, OpenExp env aenv t)] -> OpenExp env aenv t
caseof PrimBool
tag = [(PrimBool, OpenExp env aenv t)] -> OpenExp env aenv t
go
          where
            go :: [(PrimBool, OpenExp env aenv t)] -> OpenExp env aenv t
go ((PrimBool
t,OpenExp env aenv t
c):[(PrimBool, OpenExp env aenv t)]
cs)
              | PrimBool
tag PrimBool -> PrimBool -> Bool
forall a. Eq a => a -> a -> Bool
== PrimBool
t  = OpenExp env aenv t
c
              | Bool
otherwise = [(PrimBool, OpenExp env aenv t)] -> OpenExp env aenv t
go [(PrimBool, OpenExp env aenv t)]
cs
            go []
              | Just OpenExp env aenv t
d <- Maybe (OpenExp env aenv t)
def = OpenExp env aenv t
d
              | Bool
otherwise     = Format (OpenExp env aenv t) (OpenExp env aenv t)
-> OpenExp env aenv t
forall r a. HasCallStack => Format r a -> a
internalError Format (OpenExp env aenv t) (OpenExp env aenv t)
"unmatched case"

    Cond OpenExp env aenv PrimBool
c OpenExp env aenv t
t OpenExp env aenv t
e
      | PrimBool -> Bool
toBool (OpenExp env aenv PrimBool -> PrimBool
forall t'. OpenExp env aenv t' -> t'
evalE OpenExp env aenv PrimBool
c)        -> OpenExp env aenv t -> t
forall t'. OpenExp env aenv t' -> t'
evalE OpenExp env aenv t
t
      | Bool
otherwise               -> OpenExp env aenv t -> t
forall t'. OpenExp env aenv t' -> t'
evalE OpenExp env aenv t
e

    While OpenFun env aenv (t -> PrimBool)
cond OpenFun env aenv (t -> t)
body OpenExp env aenv t
seed        -> t -> t
go (OpenExp env aenv t -> t
forall t'. OpenExp env aenv t' -> t'
evalE OpenExp env aenv t
seed)
      where
        f :: t -> t
f       = OpenFun env aenv (t -> t) -> t -> t
forall f'. OpenFun env aenv f' -> f'
evalF OpenFun env aenv (t -> t)
body
        p :: t -> PrimBool
p       = OpenFun env aenv (t -> PrimBool) -> t -> PrimBool
forall f'. OpenFun env aenv f' -> f'
evalF OpenFun env aenv (t -> PrimBool)
cond
        go :: t -> t
go !t
x
          | PrimBool -> Bool
toBool (t -> PrimBool
p t
x) = t -> t
go (t -> t
f t
x)
          | Bool
otherwise    = t
x

    Index ArrayVar aenv (Array dim t)
acc OpenExp env aenv dim
ix                -> let (TupRsingle ArrayR (Array dim t)
repr, Array dim t
a) = ArrayVar aenv (Array dim t) -> (ArraysR (Array dim t), Array dim t)
forall a. ArrayVar aenv a -> WithReprs a
evalA ArrayVar aenv (Array dim t)
acc
                                   in (ArrayR (Array dim t)
repr, Array dim t
a) (ArrayR (Array dim t), Array dim t) -> dim -> t
forall sh e. (ArrayR (Array sh e), Array sh e) -> sh -> e
! OpenExp env aenv dim -> dim
forall t'. OpenExp env aenv t' -> t'
evalE OpenExp env aenv dim
ix
    LinearIndex ArrayVar aenv (Array dim t)
acc OpenExp env aenv Int
i           -> let (TupRsingle ArrayR (Array dim t)
repr, Array dim t
a) = ArrayVar aenv (Array dim t) -> (ArraysR (Array dim t), Array dim t)
forall a. ArrayVar aenv a -> WithReprs a
evalA ArrayVar aenv (Array dim t)
acc
                                       ix :: dim
ix   = ShapeR dim -> dim -> Int -> dim
forall sh. HasCallStack => ShapeR sh -> sh -> Int -> sh
fromIndex (ArrayR (Array dim t) -> ShapeR dim
forall sh e. ArrayR (Array sh e) -> ShapeR sh
arrayRshape ArrayR (Array dim t)
repr) (Array dim t -> dim
forall sh e. Array sh e -> sh
shape Array dim t
a) (OpenExp env aenv Int -> Int
forall t'. OpenExp env aenv t' -> t'
evalE OpenExp env aenv Int
i)
                                   in (ArrayR (Array dim t)
repr, Array dim t
a) (ArrayR (Array dim t), Array dim t) -> dim -> t
forall sh e. (ArrayR (Array sh e), Array sh e) -> sh -> e
! dim
ix
    Shape ArrayVar aenv (Array t e)
acc                   -> Array t e -> t
forall sh e. Array sh e -> sh
shape (Array t e -> t) -> Array t e -> t
forall a b. (a -> b) -> a -> b
$ (ArraysR (Array t e), Array t e) -> Array t e
forall a b. (a, b) -> b
snd ((ArraysR (Array t e), Array t e) -> Array t e)
-> (ArraysR (Array t e), Array t e) -> Array t e
forall a b. (a -> b) -> a -> b
$ ArrayVar aenv (Array t e) -> (ArraysR (Array t e), Array t e)
forall a. ArrayVar aenv a -> WithReprs a
evalA ArrayVar aenv (Array t e)
acc
    ShapeSize ShapeR dim
shr OpenExp env aenv dim
sh            -> ShapeR dim -> dim -> Int
forall sh. ShapeR sh -> sh -> Int
size ShapeR dim
shr (OpenExp env aenv dim -> dim
forall t'. OpenExp env aenv t' -> t'
evalE OpenExp env aenv dim
sh)
    Foreign TypeR t
_ asm (x -> t)
_ Fun () (x -> t)
f OpenExp env aenv x
e             -> Fun () (x -> t) -> Val () -> Val () -> x -> t
forall env aenv t.
HasCallStack =>
OpenFun env aenv t -> Val env -> Val aenv -> t
evalOpenFun Fun () (x -> t)
f Val ()
Empty Val ()
Empty (x -> t) -> x -> t
forall a b. (a -> b) -> a -> b
$ OpenExp env aenv x -> x
forall t'. OpenExp env aenv t' -> t'
evalE OpenExp env aenv x
e
    Coerce ScalarType a
t1 ScalarType t
t2 OpenExp env aenv a
e              -> ScalarType a -> ScalarType t -> a -> t
forall a b. ScalarType a -> ScalarType b -> a -> b
evalCoerceScalar ScalarType a
t1 ScalarType t
t2 (OpenExp env aenv a -> a
forall t'. OpenExp env aenv t' -> t'
evalE OpenExp env aenv a
e)


-- Coercions
-- ---------

-- Coercion between two scalar types. We require that the size of the source and
-- destination values are equal (this is not checked at this point).
--
evalCoerceScalar :: ScalarType a -> ScalarType b -> a -> b
evalCoerceScalar :: forall a b. ScalarType a -> ScalarType b -> a -> b
evalCoerceScalar SingleScalarType{}    SingleScalarType{} a
a = a -> b
forall a b. a -> b
unsafeCoerce a
a
evalCoerceScalar VectorScalarType{}    VectorScalarType{} a
a = a -> b
forall a b. a -> b
unsafeCoerce a
a  -- XXX: or just unpack/repack the (Vec ba#)
evalCoerceScalar (SingleScalarType SingleType a
ta) VectorScalarType{} a
a = SingleType a -> a -> Vec n a1
forall a (n :: Nat) b. SingleType a -> a -> Vec n b
vector SingleType a
ta a
a
  where
    vector :: SingleType a -> a -> Vec n b
    vector :: forall a (n :: Nat) b. SingleType a -> a -> Vec n b
vector (NumSingleType NumType a
t) = NumType a -> a -> Vec n b
forall a (n :: Nat) b. NumType a -> a -> Vec n b
num NumType a
t

    num :: NumType a -> a -> Vec n b
    num :: forall a (n :: Nat) b. NumType a -> a -> Vec n b
num (IntegralNumType IntegralType a
t) = IntegralType a -> a -> Vec n b
forall a (n :: Nat) b. IntegralType a -> a -> Vec n b
integral IntegralType a
t
    num (FloatingNumType FloatingType a
t) = FloatingType a -> a -> Vec n b
forall a (n :: Nat) b. FloatingType a -> a -> Vec n b
floating FloatingType a
t

    integral :: IntegralType a -> a -> Vec n b
    integral :: forall a (n :: Nat) b. IntegralType a -> a -> Vec n b
integral TypeInt{}     = a -> Vec n b
forall a b (n :: Nat). Prim a => a -> Vec n b
poke
    integral TypeInt8{}    = a -> Vec n b
forall a b (n :: Nat). Prim a => a -> Vec n b
poke
    integral TypeInt16{}   = a -> Vec n b
forall a b (n :: Nat). Prim a => a -> Vec n b
poke
    integral TypeInt32{}   = a -> Vec n b
forall a b (n :: Nat). Prim a => a -> Vec n b
poke
    integral TypeInt64{}   = a -> Vec n b
forall a b (n :: Nat). Prim a => a -> Vec n b
poke
    integral TypeWord{}    = a -> Vec n b
forall a b (n :: Nat). Prim a => a -> Vec n b
poke
    integral TypeWord8{}   = a -> Vec n b
forall a b (n :: Nat). Prim a => a -> Vec n b
poke
    integral TypeWord16{}  = a -> Vec n b
forall a b (n :: Nat). Prim a => a -> Vec n b
poke
    integral TypeWord32{}  = a -> Vec n b
forall a b (n :: Nat). Prim a => a -> Vec n b
poke
    integral TypeWord64{}  = a -> Vec n b
forall a b (n :: Nat). Prim a => a -> Vec n b
poke

    floating :: FloatingType a -> a -> Vec n b
    floating :: forall a (n :: Nat) b. FloatingType a -> a -> Vec n b
floating TypeHalf{}    = a -> Vec n b
forall a b (n :: Nat). Prim a => a -> Vec n b
poke
    floating TypeFloat{}   = a -> Vec n b
forall a b (n :: Nat). Prim a => a -> Vec n b
poke
    floating TypeDouble{}  = a -> Vec n b
forall a b (n :: Nat). Prim a => a -> Vec n b
poke

    {-# INLINE poke #-}
    poke :: forall a b n. Prim a => a -> Vec n b
    poke :: forall a b (n :: Nat). Prim a => a -> Vec n b
poke a
x = (forall s. ST s (Vec n b)) -> Vec n b
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Vec n b)) -> Vec n b)
-> (forall s. ST s (Vec n b)) -> Vec n b
forall a b. (a -> b) -> a -> b
$ do
      MutableByteArray s
mba <- Int -> ST s (MutableByteArray (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
newByteArray (a -> Int
forall a. Prim a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined::a))
      MutableByteArray (PrimState (ST s)) -> Int -> a -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
writeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
mba Int
0 a
x
      ByteArray ByteArray#
ba# <- MutableByteArray (PrimState (ST s)) -> ST s ByteArray
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
unsafeFreezeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
mba
      Vec n b -> ST s (Vec n b)
forall a. a -> ST s a
forall (m :: * -> *) a. Monad m => a -> m a
return (Vec n b -> ST s (Vec n b)) -> Vec n b -> ST s (Vec n b)
forall a b. (a -> b) -> a -> b
$ ByteArray# -> Vec n b
forall (n :: Nat) a. ByteArray# -> Vec n a
Vec ByteArray#
ba#

evalCoerceScalar VectorScalarType{} (SingleScalarType SingleType b
tb) a
a = SingleType b -> Vec n a1 -> b
forall b (n :: Nat) a. SingleType b -> Vec n a -> b
scalar SingleType b
tb a
Vec n a1
a
  where
    scalar :: SingleType b -> Vec n a -> b
    scalar :: forall b (n :: Nat) a. SingleType b -> Vec n a -> b
scalar (NumSingleType NumType b
t) = NumType b -> Vec n a -> b
forall b (n :: Nat) a. NumType b -> Vec n a -> b
num NumType b
t

    num :: NumType b -> Vec n a -> b
    num :: forall b (n :: Nat) a. NumType b -> Vec n a -> b
num (IntegralNumType IntegralType b
t) = IntegralType b -> Vec n a -> b
forall b (n :: Nat) a. IntegralType b -> Vec n a -> b
integral IntegralType b
t
    num (FloatingNumType FloatingType b
t) = FloatingType b -> Vec n a -> b
forall b (n :: Nat) a. FloatingType b -> Vec n a -> b
floating FloatingType b
t

    integral :: IntegralType b -> Vec n a -> b
    integral :: forall b (n :: Nat) a. IntegralType b -> Vec n a -> b
integral TypeInt{}     = Vec n a -> b
forall a (n :: Nat) b. Prim a => Vec n b -> a
peek
    integral TypeInt8{}    = Vec n a -> b
forall a (n :: Nat) b. Prim a => Vec n b -> a
peek
    integral TypeInt16{}   = Vec n a -> b
forall a (n :: Nat) b. Prim a => Vec n b -> a
peek
    integral TypeInt32{}   = Vec n a -> b
forall a (n :: Nat) b. Prim a => Vec n b -> a
peek
    integral TypeInt64{}   = Vec n a -> b
forall a (n :: Nat) b. Prim a => Vec n b -> a
peek
    integral TypeWord{}    = Vec n a -> b
forall a (n :: Nat) b. Prim a => Vec n b -> a
peek
    integral TypeWord8{}   = Vec n a -> b
forall a (n :: Nat) b. Prim a => Vec n b -> a
peek
    integral TypeWord16{}  = Vec n a -> b
forall a (n :: Nat) b. Prim a => Vec n b -> a
peek
    integral TypeWord32{}  = Vec n a -> b
forall a (n :: Nat) b. Prim a => Vec n b -> a
peek
    integral TypeWord64{}  = Vec n a -> b
forall a (n :: Nat) b. Prim a => Vec n b -> a
peek

    floating :: FloatingType b -> Vec n a -> b
    floating :: forall b (n :: Nat) a. FloatingType b -> Vec n a -> b
floating TypeHalf{}    = Vec n a -> b
forall a (n :: Nat) b. Prim a => Vec n b -> a
peek
    floating TypeFloat{}   = Vec n a -> b
forall a (n :: Nat) b. Prim a => Vec n b -> a
peek
    floating TypeDouble{}  = Vec n a -> b
forall a (n :: Nat) b. Prim a => Vec n b -> a
peek

    {-# INLINE peek #-}
    peek :: Prim a => Vec n b -> a
    peek :: forall a (n :: Nat) b. Prim a => Vec n b -> a
peek (Vec ByteArray#
ba#) = ByteArray -> Int -> a
forall a. Prim a => ByteArray -> Int -> a
indexByteArray (ByteArray# -> ByteArray
ByteArray ByteArray#
ba#) Int
0


-- Scalar primitives
-- -----------------

evalPrimConst :: PrimConst a -> a
evalPrimConst :: forall a. PrimConst a -> a
evalPrimConst (PrimMinBound BoundedType a
ty) = BoundedType a -> a
forall a. BoundedType a -> a
evalMinBound BoundedType a
ty
evalPrimConst (PrimMaxBound BoundedType a
ty) = BoundedType a -> a
forall a. BoundedType a -> a
evalMaxBound BoundedType a
ty
evalPrimConst (PrimPi       FloatingType a
ty) = FloatingType a -> a
forall a. FloatingType a -> a
evalPi FloatingType a
ty

evalPrim :: PrimFun (a -> r) -> (a -> r)
evalPrim :: forall a r. PrimFun (a -> r) -> a -> r
evalPrim (PrimAdd                NumType a
ty) = NumType r -> (r, r) -> r
forall a. NumType a -> (a, a) -> a
evalAdd NumType r
NumType a
ty
evalPrim (PrimSub                NumType a
ty) = NumType r -> (r, r) -> r
forall a. NumType a -> (a, a) -> a
evalSub NumType r
NumType a
ty
evalPrim (PrimMul                NumType a
ty) = NumType r -> (r, r) -> r
forall a. NumType a -> (a, a) -> a
evalMul NumType r
NumType a
ty
evalPrim (PrimNeg                NumType a
ty) = NumType a -> a -> a
forall a. NumType a -> a -> a
evalNeg NumType a
NumType a
ty
evalPrim (PrimAbs                NumType a
ty) = NumType a -> a -> a
forall a. NumType a -> a -> a
evalAbs NumType a
NumType a
ty
evalPrim (PrimSig                NumType a
ty) = NumType a -> a -> a
forall a. NumType a -> a -> a
evalSig NumType a
NumType a
ty
evalPrim (PrimQuot               IntegralType a
ty) = IntegralType r -> (r, r) -> r
forall a. IntegralType a -> (a, a) -> a
evalQuot IntegralType r
IntegralType a
ty
evalPrim (PrimRem                IntegralType a
ty) = IntegralType r -> (r, r) -> r
forall a. IntegralType a -> (a, a) -> a
evalRem IntegralType r
IntegralType a
ty
evalPrim (PrimQuotRem            IntegralType a
ty) = IntegralType a -> (a, a) -> (a, a)
forall a. IntegralType a -> (a, a) -> (a, a)
evalQuotRem IntegralType a
ty
evalPrim (PrimIDiv               IntegralType a
ty) = IntegralType r -> (r, r) -> r
forall a. IntegralType a -> (a, a) -> a
evalIDiv IntegralType r
IntegralType a
ty
evalPrim (PrimMod                IntegralType a
ty) = IntegralType r -> (r, r) -> r
forall a. IntegralType a -> (a, a) -> a
evalMod IntegralType r
IntegralType a
ty
evalPrim (PrimDivMod             IntegralType a
ty) = IntegralType a -> (a, a) -> (a, a)
forall a. IntegralType a -> (a, a) -> (a, a)
evalDivMod IntegralType a
ty
evalPrim (PrimBAnd               IntegralType a
ty) = IntegralType r -> (r, r) -> r
forall a. IntegralType a -> (a, a) -> a
evalBAnd IntegralType r
IntegralType a
ty
evalPrim (PrimBOr                IntegralType a
ty) = IntegralType r -> (r, r) -> r
forall a. IntegralType a -> (a, a) -> a
evalBOr IntegralType r
IntegralType a
ty
evalPrim (PrimBXor               IntegralType a
ty) = IntegralType r -> (r, r) -> r
forall a. IntegralType a -> (a, a) -> a
evalBXor IntegralType r
IntegralType a
ty
evalPrim (PrimBNot               IntegralType a
ty) = IntegralType a -> a -> a
forall a. IntegralType a -> a -> a
evalBNot IntegralType a
IntegralType a
ty
evalPrim (PrimBShiftL            IntegralType a
ty) = IntegralType r -> (r, Int) -> r
forall a. IntegralType a -> (a, Int) -> a
evalBShiftL IntegralType r
IntegralType a
ty
evalPrim (PrimBShiftR            IntegralType a
ty) = IntegralType r -> (r, Int) -> r
forall a. IntegralType a -> (a, Int) -> a
evalBShiftR IntegralType r
IntegralType a
ty
evalPrim (PrimBRotateL           IntegralType a
ty) = IntegralType r -> (r, Int) -> r
forall a. IntegralType a -> (a, Int) -> a
evalBRotateL IntegralType r
IntegralType a
ty
evalPrim (PrimBRotateR           IntegralType a
ty) = IntegralType r -> (r, Int) -> r
forall a. IntegralType a -> (a, Int) -> a
evalBRotateR IntegralType r
IntegralType a
ty
evalPrim (PrimPopCount           IntegralType a
ty) = IntegralType a -> a -> Int
forall a. IntegralType a -> a -> Int
evalPopCount IntegralType a
IntegralType a
ty
evalPrim (PrimCountLeadingZeros  IntegralType a
ty) = IntegralType a -> a -> Int
forall a. IntegralType a -> a -> Int
evalCountLeadingZeros IntegralType a
IntegralType a
ty
evalPrim (PrimCountTrailingZeros IntegralType a
ty) = IntegralType a -> a -> Int
forall a. IntegralType a -> a -> Int
evalCountTrailingZeros IntegralType a
IntegralType a
ty
evalPrim (PrimFDiv               FloatingType a
ty) = FloatingType r -> (r, r) -> r
forall a. FloatingType a -> (a, a) -> a
evalFDiv FloatingType r
FloatingType a
ty
evalPrim (PrimRecip              FloatingType a
ty) = FloatingType a -> a -> a
forall a. FloatingType a -> a -> a
evalRecip FloatingType a
FloatingType a
ty
evalPrim (PrimSin                FloatingType a
ty) = FloatingType a -> a -> a
forall a. FloatingType a -> a -> a
evalSin FloatingType a
FloatingType a
ty
evalPrim (PrimCos                FloatingType a
ty) = FloatingType a -> a -> a
forall a. FloatingType a -> a -> a
evalCos FloatingType a
FloatingType a
ty
evalPrim (PrimTan                FloatingType a
ty) = FloatingType a -> a -> a
forall a. FloatingType a -> a -> a
evalTan FloatingType a
FloatingType a
ty
evalPrim (PrimAsin               FloatingType a
ty) = FloatingType a -> a -> a
forall a. FloatingType a -> a -> a
evalAsin FloatingType a
FloatingType a
ty
evalPrim (PrimAcos               FloatingType a
ty) = FloatingType a -> a -> a
forall a. FloatingType a -> a -> a
evalAcos FloatingType a
FloatingType a
ty
evalPrim (PrimAtan               FloatingType a
ty) = FloatingType a -> a -> a
forall a. FloatingType a -> a -> a
evalAtan FloatingType a
FloatingType a
ty
evalPrim (PrimSinh               FloatingType a
ty) = FloatingType a -> a -> a
forall a. FloatingType a -> a -> a
evalSinh FloatingType a
FloatingType a
ty
evalPrim (PrimCosh               FloatingType a
ty) = FloatingType a -> a -> a
forall a. FloatingType a -> a -> a
evalCosh FloatingType a
FloatingType a
ty
evalPrim (PrimTanh               FloatingType a
ty) = FloatingType a -> a -> a
forall a. FloatingType a -> a -> a
evalTanh FloatingType a
FloatingType a
ty
evalPrim (PrimAsinh              FloatingType a
ty) = FloatingType a -> a -> a
forall a. FloatingType a -> a -> a
evalAsinh FloatingType a
FloatingType a
ty
evalPrim (PrimAcosh              FloatingType a
ty) = FloatingType a -> a -> a
forall a. FloatingType a -> a -> a
evalAcosh FloatingType a
FloatingType a
ty
evalPrim (PrimAtanh              FloatingType a
ty) = FloatingType a -> a -> a
forall a. FloatingType a -> a -> a
evalAtanh FloatingType a
FloatingType a
ty
evalPrim (PrimExpFloating        FloatingType a
ty) = FloatingType a -> a -> a
forall a. FloatingType a -> a -> a
evalExpFloating FloatingType a
FloatingType a
ty
evalPrim (PrimSqrt               FloatingType a
ty) = FloatingType a -> a -> a
forall a. FloatingType a -> a -> a
evalSqrt FloatingType a
FloatingType a
ty
evalPrim (PrimLog                FloatingType a
ty) = FloatingType a -> a -> a
forall a. FloatingType a -> a -> a
evalLog FloatingType a
FloatingType a
ty
evalPrim (PrimFPow               FloatingType a
ty) = FloatingType r -> (r, r) -> r
forall a. FloatingType a -> (a, a) -> a
evalFPow FloatingType r
FloatingType a
ty
evalPrim (PrimLogBase            FloatingType a
ty) = FloatingType r -> (r, r) -> r
forall a. FloatingType a -> (a, a) -> a
evalLogBase FloatingType r
FloatingType a
ty
evalPrim (PrimTruncate        FloatingType a
ta IntegralType b
tb) = FloatingType a -> IntegralType r -> a -> r
forall a b. FloatingType a -> IntegralType b -> a -> b
evalTruncate FloatingType a
FloatingType a
ta IntegralType r
IntegralType b
tb
evalPrim (PrimRound           FloatingType a
ta IntegralType b
tb) = FloatingType a -> IntegralType r -> a -> r
forall a b. FloatingType a -> IntegralType b -> a -> b
evalRound FloatingType a
FloatingType a
ta IntegralType r
IntegralType b
tb
evalPrim (PrimFloor           FloatingType a
ta IntegralType b
tb) = FloatingType a -> IntegralType r -> a -> r
forall a b. FloatingType a -> IntegralType b -> a -> b
evalFloor FloatingType a
FloatingType a
ta IntegralType r
IntegralType b
tb
evalPrim (PrimCeiling         FloatingType a
ta IntegralType b
tb) = FloatingType a -> IntegralType r -> a -> r
forall a b. FloatingType a -> IntegralType b -> a -> b
evalCeiling FloatingType a
FloatingType a
ta IntegralType r
IntegralType b
tb
evalPrim (PrimAtan2              FloatingType a
ty) = FloatingType r -> (r, r) -> r
forall a. FloatingType a -> (a, a) -> a
evalAtan2 FloatingType r
FloatingType a
ty
evalPrim (PrimIsNaN              FloatingType a
ty) = FloatingType a -> a -> PrimBool
forall a. FloatingType a -> a -> PrimBool
evalIsNaN FloatingType a
FloatingType a
ty
evalPrim (PrimIsInfinite         FloatingType a
ty) = FloatingType a -> a -> PrimBool
forall a. FloatingType a -> a -> PrimBool
evalIsInfinite FloatingType a
FloatingType a
ty
evalPrim (PrimLt                 SingleType a
ty) = SingleType a -> (a, a) -> PrimBool
forall a. SingleType a -> (a, a) -> PrimBool
evalLt SingleType a
ty
evalPrim (PrimGt                 SingleType a
ty) = SingleType a -> (a, a) -> PrimBool
forall a. SingleType a -> (a, a) -> PrimBool
evalGt SingleType a
ty
evalPrim (PrimLtEq               SingleType a
ty) = SingleType a -> (a, a) -> PrimBool
forall a. SingleType a -> (a, a) -> PrimBool
evalLtEq SingleType a
ty
evalPrim (PrimGtEq               SingleType a
ty) = SingleType a -> (a, a) -> PrimBool
forall a. SingleType a -> (a, a) -> PrimBool
evalGtEq SingleType a
ty
evalPrim (PrimEq                 SingleType a
ty) = SingleType a -> (a, a) -> PrimBool
forall a. SingleType a -> (a, a) -> PrimBool
evalEq SingleType a
ty
evalPrim (PrimNEq                SingleType a
ty) = SingleType a -> (a, a) -> PrimBool
forall a. SingleType a -> (a, a) -> PrimBool
evalNEq SingleType a
ty
evalPrim (PrimMax                SingleType a
ty) = SingleType r -> (r, r) -> r
forall a. SingleType a -> (a, a) -> a
evalMax SingleType r
SingleType a
ty
evalPrim (PrimMin                SingleType a
ty) = SingleType r -> (r, r) -> r
forall a. SingleType a -> (a, a) -> a
evalMin SingleType r
SingleType a
ty
evalPrim PrimFun (a -> r)
PrimLAnd                    = a -> r
(PrimBool, PrimBool) -> PrimBool
evalLAnd
evalPrim PrimFun (a -> r)
PrimLOr                     = a -> r
(PrimBool, PrimBool) -> PrimBool
evalLOr
evalPrim PrimFun (a -> r)
PrimLNot                    = a -> r
PrimBool -> PrimBool
evalLNot
evalPrim (PrimFromIntegral IntegralType a
ta NumType b
tb)    = IntegralType a -> NumType r -> a -> r
forall a b. IntegralType a -> NumType b -> a -> b
evalFromIntegral IntegralType a
IntegralType a
ta NumType r
NumType b
tb
evalPrim (PrimToFloating NumType a
ta FloatingType b
tb)      = NumType a -> FloatingType r -> a -> r
forall a b. NumType a -> FloatingType b -> a -> b
evalToFloating NumType a
NumType a
ta FloatingType r
FloatingType b
tb


-- Implementation of scalar primitives
-- -----------------------------------

toBool :: PrimBool -> Bool
toBool :: PrimBool -> Bool
toBool PrimBool
0 = Bool
False
toBool PrimBool
_ = Bool
True

fromBool :: Bool -> PrimBool
fromBool :: Bool -> PrimBool
fromBool Bool
False = PrimBool
0
fromBool Bool
True  = PrimBool
1

evalLAnd :: (PrimBool, PrimBool) -> PrimBool
evalLAnd :: (PrimBool, PrimBool) -> PrimBool
evalLAnd (PrimBool
x, PrimBool
y) = Bool -> PrimBool
fromBool (PrimBool -> Bool
toBool PrimBool
x Bool -> Bool -> Bool
&& PrimBool -> Bool
toBool PrimBool
y)

evalLOr  :: (PrimBool, PrimBool) -> PrimBool
evalLOr :: (PrimBool, PrimBool) -> PrimBool
evalLOr (PrimBool
x, PrimBool
y) = Bool -> PrimBool
fromBool (PrimBool -> Bool
toBool PrimBool
x Bool -> Bool -> Bool
|| PrimBool -> Bool
toBool PrimBool
y)

evalLNot :: PrimBool -> PrimBool
evalLNot :: PrimBool -> PrimBool
evalLNot = Bool -> PrimBool
fromBool (Bool -> PrimBool) -> (PrimBool -> Bool) -> PrimBool -> PrimBool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> Bool
not (Bool -> Bool) -> (PrimBool -> Bool) -> PrimBool -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PrimBool -> Bool
toBool

evalFromIntegral :: IntegralType a -> NumType b -> a -> b
evalFromIntegral :: forall a b. IntegralType a -> NumType b -> a -> b
evalFromIntegral IntegralType a
ta (IntegralNumType IntegralType b
tb)
  | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ta
  , IntegralDict b
IntegralDict <- IntegralType b -> IntegralDict b
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType b
tb
  = a -> b
forall a b. (Integral a, Num b) => a -> b
fromIntegral

evalFromIntegral IntegralType a
ta (FloatingNumType FloatingType b
tb)
  | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ta
  , FloatingDict b
FloatingDict <- FloatingType b -> FloatingDict b
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType b
tb
  = a -> b
forall a b. (Integral a, Num b) => a -> b
fromIntegral

evalToFloating :: NumType a -> FloatingType b -> a -> b
evalToFloating :: forall a b. NumType a -> FloatingType b -> a -> b
evalToFloating (IntegralNumType IntegralType a
ta) FloatingType b
tb
  | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ta
  , FloatingDict b
FloatingDict <- FloatingType b -> FloatingDict b
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType b
tb
  = a -> b
forall a b. (Real a, Fractional b) => a -> b
realToFrac

evalToFloating (FloatingNumType FloatingType a
ta) FloatingType b
tb
  | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ta
  , FloatingDict b
FloatingDict <- FloatingType b -> FloatingDict b
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType b
tb
  = a -> b
forall a b. (Real a, Fractional b) => a -> b
realToFrac


-- Extract methods from reified dictionaries
--

-- Constant methods of Bounded
--

evalMinBound :: BoundedType a -> a
evalMinBound :: forall a. BoundedType a -> a
evalMinBound (IntegralBoundedType IntegralType a
ty)
  | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty
  = a
forall a. Bounded a => a
minBound

evalMaxBound :: BoundedType a -> a
evalMaxBound :: forall a. BoundedType a -> a
evalMaxBound (IntegralBoundedType IntegralType a
ty)
  | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty
  = a
forall a. Bounded a => a
maxBound

-- Constant method of floating
--

evalPi :: FloatingType a -> a
evalPi :: forall a. FloatingType a -> a
evalPi FloatingType a
ty | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = a
forall a. Floating a => a
pi

evalSin :: FloatingType a -> (a -> a)
evalSin :: forall a. FloatingType a -> a -> a
evalSin FloatingType a
ty | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = a -> a
forall a. Floating a => a -> a
sin

evalCos :: FloatingType a -> (a -> a)
evalCos :: forall a. FloatingType a -> a -> a
evalCos FloatingType a
ty | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = a -> a
forall a. Floating a => a -> a
cos

evalTan :: FloatingType a -> (a -> a)
evalTan :: forall a. FloatingType a -> a -> a
evalTan FloatingType a
ty | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = a -> a
forall a. Floating a => a -> a
tan

evalAsin :: FloatingType a -> (a -> a)
evalAsin :: forall a. FloatingType a -> a -> a
evalAsin FloatingType a
ty | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = a -> a
forall a. Floating a => a -> a
asin

evalAcos :: FloatingType a -> (a -> a)
evalAcos :: forall a. FloatingType a -> a -> a
evalAcos FloatingType a
ty | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = a -> a
forall a. Floating a => a -> a
acos

evalAtan :: FloatingType a -> (a -> a)
evalAtan :: forall a. FloatingType a -> a -> a
evalAtan FloatingType a
ty | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = a -> a
forall a. Floating a => a -> a
atan

evalSinh :: FloatingType a -> (a -> a)
evalSinh :: forall a. FloatingType a -> a -> a
evalSinh FloatingType a
ty | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = a -> a
forall a. Floating a => a -> a
sinh

evalCosh :: FloatingType a -> (a -> a)
evalCosh :: forall a. FloatingType a -> a -> a
evalCosh FloatingType a
ty | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = a -> a
forall a. Floating a => a -> a
cosh

evalTanh :: FloatingType a -> (a -> a)
evalTanh :: forall a. FloatingType a -> a -> a
evalTanh FloatingType a
ty | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = a -> a
forall a. Floating a => a -> a
tanh

evalAsinh :: FloatingType a -> (a -> a)
evalAsinh :: forall a. FloatingType a -> a -> a
evalAsinh FloatingType a
ty | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = a -> a
forall a. Floating a => a -> a
asinh

evalAcosh :: FloatingType a -> (a -> a)
evalAcosh :: forall a. FloatingType a -> a -> a
evalAcosh FloatingType a
ty | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = a -> a
forall a. Floating a => a -> a
acosh

evalAtanh :: FloatingType a -> (a -> a)
evalAtanh :: forall a. FloatingType a -> a -> a
evalAtanh FloatingType a
ty | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = a -> a
forall a. Floating a => a -> a
atanh

evalExpFloating :: FloatingType a -> (a -> a)
evalExpFloating :: forall a. FloatingType a -> a -> a
evalExpFloating FloatingType a
ty | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = a -> a
forall a. Floating a => a -> a
exp

evalSqrt :: FloatingType a -> (a -> a)
evalSqrt :: forall a. FloatingType a -> a -> a
evalSqrt FloatingType a
ty | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = a -> a
forall a. Floating a => a -> a
sqrt

evalLog :: FloatingType a -> (a -> a)
evalLog :: forall a. FloatingType a -> a -> a
evalLog FloatingType a
ty | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = a -> a
forall a. Floating a => a -> a
log

evalFPow :: FloatingType a -> ((a, a) -> a)
evalFPow :: forall a. FloatingType a -> (a, a) -> a
evalFPow FloatingType a
ty | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = (a -> a -> a) -> (a, a) -> a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> a
forall a. Floating a => a -> a -> a
(**)

evalLogBase :: FloatingType a -> ((a, a) -> a)
evalLogBase :: forall a. FloatingType a -> (a, a) -> a
evalLogBase FloatingType a
ty | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = (a -> a -> a) -> (a, a) -> a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> a
forall a. Floating a => a -> a -> a
logBase

evalTruncate :: FloatingType a -> IntegralType b -> (a -> b)
evalTruncate :: forall a b. FloatingType a -> IntegralType b -> a -> b
evalTruncate FloatingType a
ta IntegralType b
tb
  | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ta
  , IntegralDict b
IntegralDict <- IntegralType b -> IntegralDict b
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType b
tb
  = a -> b
forall b. Integral b => a -> b
forall a b. (RealFrac a, Integral b) => a -> b
truncate

evalRound :: FloatingType a -> IntegralType b -> (a -> b)
evalRound :: forall a b. FloatingType a -> IntegralType b -> a -> b
evalRound FloatingType a
ta IntegralType b
tb
  | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ta
  , IntegralDict b
IntegralDict <- IntegralType b -> IntegralDict b
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType b
tb
  = a -> b
forall b. Integral b => a -> b
forall a b. (RealFrac a, Integral b) => a -> b
round

evalFloor :: FloatingType a -> IntegralType b -> (a -> b)
evalFloor :: forall a b. FloatingType a -> IntegralType b -> a -> b
evalFloor FloatingType a
ta IntegralType b
tb
  | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ta
  , IntegralDict b
IntegralDict <- IntegralType b -> IntegralDict b
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType b
tb
  = a -> b
forall b. Integral b => a -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor

evalCeiling :: FloatingType a -> IntegralType b -> (a -> b)
evalCeiling :: forall a b. FloatingType a -> IntegralType b -> a -> b
evalCeiling FloatingType a
ta IntegralType b
tb
  | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ta
  , IntegralDict b
IntegralDict <- IntegralType b -> IntegralDict b
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType b
tb
  = a -> b
forall b. Integral b => a -> b
forall a b. (RealFrac a, Integral b) => a -> b
ceiling

evalAtan2 :: FloatingType a -> ((a, a) -> a)
evalAtan2 :: forall a. FloatingType a -> (a, a) -> a
evalAtan2 FloatingType a
ty | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = (a -> a -> a) -> (a, a) -> a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> a
forall a. RealFloat a => a -> a -> a
atan2

evalIsNaN :: FloatingType a -> (a -> PrimBool)
evalIsNaN :: forall a. FloatingType a -> a -> PrimBool
evalIsNaN FloatingType a
ty | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = Bool -> PrimBool
fromBool (Bool -> PrimBool) -> (a -> Bool) -> a -> PrimBool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Bool
forall a. RealFloat a => a -> Bool
isNaN

evalIsInfinite :: FloatingType a -> (a -> PrimBool)
evalIsInfinite :: forall a. FloatingType a -> a -> PrimBool
evalIsInfinite FloatingType a
ty | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = Bool -> PrimBool
fromBool (Bool -> PrimBool) -> (a -> Bool) -> a -> PrimBool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Bool
forall a. RealFloat a => a -> Bool
isInfinite


-- Methods of Num
--

evalAdd :: NumType a -> ((a, a) -> a)
evalAdd :: forall a. NumType a -> (a, a) -> a
evalAdd (IntegralNumType IntegralType a
ty) | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty = (a -> a -> a) -> (a, a) -> a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> a
forall a. Num a => a -> a -> a
(+)
evalAdd (FloatingNumType FloatingType a
ty) | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = (a -> a -> a) -> (a, a) -> a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> a
forall a. Num a => a -> a -> a
(+)

evalSub :: NumType a -> ((a, a) -> a)
evalSub :: forall a. NumType a -> (a, a) -> a
evalSub (IntegralNumType IntegralType a
ty) | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty = (a -> a -> a) -> (a, a) -> a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (-)
evalSub (FloatingNumType FloatingType a
ty) | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = (a -> a -> a) -> (a, a) -> a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (-)

evalMul :: NumType a -> ((a, a) -> a)
evalMul :: forall a. NumType a -> (a, a) -> a
evalMul (IntegralNumType IntegralType a
ty) | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty = (a -> a -> a) -> (a, a) -> a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> a
forall a. Num a => a -> a -> a
(*)
evalMul (FloatingNumType FloatingType a
ty) | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = (a -> a -> a) -> (a, a) -> a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> a
forall a. Num a => a -> a -> a
(*)

evalNeg :: NumType a -> (a -> a)
evalNeg :: forall a. NumType a -> a -> a
evalNeg (IntegralNumType IntegralType a
ty) | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty = a -> a
forall a. Num a => a -> a
negate
evalNeg (FloatingNumType FloatingType a
ty) | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = a -> a
forall a. Num a => a -> a
negate

evalAbs :: NumType a -> (a -> a)
evalAbs :: forall a. NumType a -> a -> a
evalAbs (IntegralNumType IntegralType a
ty) | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty = a -> a
forall a. Num a => a -> a
abs
evalAbs (FloatingNumType FloatingType a
ty) | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = a -> a
forall a. Num a => a -> a
abs

evalSig :: NumType a -> (a -> a)
evalSig :: forall a. NumType a -> a -> a
evalSig (IntegralNumType IntegralType a
ty) | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty = a -> a
forall a. Num a => a -> a
signum
evalSig (FloatingNumType FloatingType a
ty) | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = a -> a
forall a. Num a => a -> a
signum

evalQuot :: IntegralType a -> ((a, a) -> a)
evalQuot :: forall a. IntegralType a -> (a, a) -> a
evalQuot IntegralType a
ty | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty = (a -> a -> a) -> (a, a) -> a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> a
forall a. Integral a => a -> a -> a
quot

evalRem :: IntegralType a -> ((a, a) -> a)
evalRem :: forall a. IntegralType a -> (a, a) -> a
evalRem IntegralType a
ty | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty = (a -> a -> a) -> (a, a) -> a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> a
forall a. Integral a => a -> a -> a
rem

evalQuotRem :: IntegralType a -> ((a, a) -> (a, a))
evalQuotRem :: forall a. IntegralType a -> (a, a) -> (a, a)
evalQuotRem IntegralType a
ty | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty = (a -> a -> (a, a)) -> (a, a) -> (a, a)
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> (a, a)
forall a. Integral a => a -> a -> (a, a)
quotRem

evalIDiv :: IntegralType a -> ((a, a) -> a)
evalIDiv :: forall a. IntegralType a -> (a, a) -> a
evalIDiv IntegralType a
ty | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty = (a -> a -> a) -> (a, a) -> a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> a
forall a. Integral a => a -> a -> a
div

evalMod :: IntegralType a -> ((a, a) -> a)
evalMod :: forall a. IntegralType a -> (a, a) -> a
evalMod IntegralType a
ty | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty = (a -> a -> a) -> (a, a) -> a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> a
forall a. Integral a => a -> a -> a
mod

evalDivMod :: IntegralType a -> ((a, a) -> (a, a))
evalDivMod :: forall a. IntegralType a -> (a, a) -> (a, a)
evalDivMod IntegralType a
ty | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty = (a -> a -> (a, a)) -> (a, a) -> (a, a)
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> (a, a)
forall a. Integral a => a -> a -> (a, a)
divMod

evalBAnd :: IntegralType a -> ((a, a) -> a)
evalBAnd :: forall a. IntegralType a -> (a, a) -> a
evalBAnd IntegralType a
ty | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty = (a -> a -> a) -> (a, a) -> a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> a
forall a. Bits a => a -> a -> a
(.&.)

evalBOr :: IntegralType a -> ((a, a) -> a)
evalBOr :: forall a. IntegralType a -> (a, a) -> a
evalBOr IntegralType a
ty | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty = (a -> a -> a) -> (a, a) -> a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> a
forall a. Bits a => a -> a -> a
(.|.)

evalBXor :: IntegralType a -> ((a, a) -> a)
evalBXor :: forall a. IntegralType a -> (a, a) -> a
evalBXor IntegralType a
ty | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty = (a -> a -> a) -> (a, a) -> a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> a
forall a. Bits a => a -> a -> a
xor

evalBNot :: IntegralType a -> (a -> a)
evalBNot :: forall a. IntegralType a -> a -> a
evalBNot IntegralType a
ty | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty = a -> a
forall a. Bits a => a -> a
complement

evalBShiftL :: IntegralType a -> ((a, Int) -> a)
evalBShiftL :: forall a. IntegralType a -> (a, Int) -> a
evalBShiftL IntegralType a
ty | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty = (a -> Int -> a) -> (a, Int) -> a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> Int -> a
forall a. Bits a => a -> Int -> a
shiftL

evalBShiftR :: IntegralType a -> ((a, Int) -> a)
evalBShiftR :: forall a. IntegralType a -> (a, Int) -> a
evalBShiftR IntegralType a
ty | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty = (a -> Int -> a) -> (a, Int) -> a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> Int -> a
forall a. Bits a => a -> Int -> a
shiftR

evalBRotateL :: IntegralType a -> ((a, Int) -> a)
evalBRotateL :: forall a. IntegralType a -> (a, Int) -> a
evalBRotateL IntegralType a
ty | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty = (a -> Int -> a) -> (a, Int) -> a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> Int -> a
forall a. Bits a => a -> Int -> a
rotateL

evalBRotateR :: IntegralType a -> ((a, Int) -> a)
evalBRotateR :: forall a. IntegralType a -> (a, Int) -> a
evalBRotateR IntegralType a
ty | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty = (a -> Int -> a) -> (a, Int) -> a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> Int -> a
forall a. Bits a => a -> Int -> a
rotateR

evalPopCount :: IntegralType a -> (a -> Int)
evalPopCount :: forall a. IntegralType a -> a -> Int
evalPopCount IntegralType a
ty | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty = a -> Int
forall a. Bits a => a -> Int
popCount

evalCountLeadingZeros :: IntegralType a -> (a -> Int)
evalCountLeadingZeros :: forall a. IntegralType a -> a -> Int
evalCountLeadingZeros IntegralType a
ty | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty = a -> Int
forall b. FiniteBits b => b -> Int
countLeadingZeros

evalCountTrailingZeros :: IntegralType a -> (a -> Int)
evalCountTrailingZeros :: forall a. IntegralType a -> a -> Int
evalCountTrailingZeros IntegralType a
ty | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty = a -> Int
forall b. FiniteBits b => b -> Int
countTrailingZeros

evalFDiv :: FloatingType a -> ((a, a) -> a)
evalFDiv :: forall a. FloatingType a -> (a, a) -> a
evalFDiv FloatingType a
ty | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = (a -> a -> a) -> (a, a) -> a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> a
forall a. Fractional a => a -> a -> a
(/)

evalRecip :: FloatingType a -> (a -> a)
evalRecip :: forall a. FloatingType a -> a -> a
evalRecip FloatingType a
ty | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = a -> a
forall a. Fractional a => a -> a
recip


evalLt :: SingleType a -> ((a, a) -> PrimBool)
evalLt :: forall a. SingleType a -> (a, a) -> PrimBool
evalLt (NumSingleType (IntegralNumType IntegralType a
ty)) | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty = Bool -> PrimBool
fromBool (Bool -> PrimBool) -> ((a, a) -> Bool) -> (a, a) -> PrimBool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> Bool) -> (a, a) -> Bool
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(<)
evalLt (NumSingleType (FloatingNumType FloatingType a
ty)) | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = Bool -> PrimBool
fromBool (Bool -> PrimBool) -> ((a, a) -> Bool) -> (a, a) -> PrimBool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> Bool) -> (a, a) -> Bool
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(<)

evalGt :: SingleType a -> ((a, a) -> PrimBool)
evalGt :: forall a. SingleType a -> (a, a) -> PrimBool
evalGt (NumSingleType (IntegralNumType IntegralType a
ty)) | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty = Bool -> PrimBool
fromBool (Bool -> PrimBool) -> ((a, a) -> Bool) -> (a, a) -> PrimBool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> Bool) -> (a, a) -> Bool
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(>)
evalGt (NumSingleType (FloatingNumType FloatingType a
ty)) | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = Bool -> PrimBool
fromBool (Bool -> PrimBool) -> ((a, a) -> Bool) -> (a, a) -> PrimBool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> Bool) -> (a, a) -> Bool
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(>)

evalLtEq :: SingleType a -> ((a, a) -> PrimBool)
evalLtEq :: forall a. SingleType a -> (a, a) -> PrimBool
evalLtEq (NumSingleType (IntegralNumType IntegralType a
ty)) | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty = Bool -> PrimBool
fromBool (Bool -> PrimBool) -> ((a, a) -> Bool) -> (a, a) -> PrimBool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> Bool) -> (a, a) -> Bool
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(<=)
evalLtEq (NumSingleType (FloatingNumType FloatingType a
ty)) | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = Bool -> PrimBool
fromBool (Bool -> PrimBool) -> ((a, a) -> Bool) -> (a, a) -> PrimBool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> Bool) -> (a, a) -> Bool
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(<=)

evalGtEq :: SingleType a -> ((a, a) -> PrimBool)
evalGtEq :: forall a. SingleType a -> (a, a) -> PrimBool
evalGtEq (NumSingleType (IntegralNumType IntegralType a
ty)) | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty = Bool -> PrimBool
fromBool (Bool -> PrimBool) -> ((a, a) -> Bool) -> (a, a) -> PrimBool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> Bool) -> (a, a) -> Bool
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(>=)
evalGtEq (NumSingleType (FloatingNumType FloatingType a
ty)) | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = Bool -> PrimBool
fromBool (Bool -> PrimBool) -> ((a, a) -> Bool) -> (a, a) -> PrimBool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> Bool) -> (a, a) -> Bool
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> Bool
forall a. Ord a => a -> a -> Bool
(>=)

evalEq :: SingleType a -> ((a, a) -> PrimBool)
evalEq :: forall a. SingleType a -> (a, a) -> PrimBool
evalEq (NumSingleType (IntegralNumType IntegralType a
ty)) | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty = Bool -> PrimBool
fromBool (Bool -> PrimBool) -> ((a, a) -> Bool) -> (a, a) -> PrimBool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> Bool) -> (a, a) -> Bool
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> Bool
forall a. Eq a => a -> a -> Bool
(==)
evalEq (NumSingleType (FloatingNumType FloatingType a
ty)) | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = Bool -> PrimBool
fromBool (Bool -> PrimBool) -> ((a, a) -> Bool) -> (a, a) -> PrimBool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> Bool) -> (a, a) -> Bool
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> Bool
forall a. Eq a => a -> a -> Bool
(==)

evalNEq :: SingleType a -> ((a, a) -> PrimBool)
evalNEq :: forall a. SingleType a -> (a, a) -> PrimBool
evalNEq (NumSingleType (IntegralNumType IntegralType a
ty)) | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty = Bool -> PrimBool
fromBool (Bool -> PrimBool) -> ((a, a) -> Bool) -> (a, a) -> PrimBool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> Bool) -> (a, a) -> Bool
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> Bool
forall a. Eq a => a -> a -> Bool
(/=)
evalNEq (NumSingleType (FloatingNumType FloatingType a
ty)) | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = Bool -> PrimBool
fromBool (Bool -> PrimBool) -> ((a, a) -> Bool) -> (a, a) -> PrimBool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> Bool) -> (a, a) -> Bool
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> Bool
forall a. Eq a => a -> a -> Bool
(/=)

evalMax :: SingleType a -> ((a, a) -> a)
evalMax :: forall a. SingleType a -> (a, a) -> a
evalMax (NumSingleType (IntegralNumType IntegralType a
ty)) | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty = (a -> a -> a) -> (a, a) -> a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> a
forall a. Ord a => a -> a -> a
max
evalMax (NumSingleType (FloatingNumType FloatingType a
ty)) | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = (a -> a -> a) -> (a, a) -> a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> a
forall a. Ord a => a -> a -> a
max

evalMin :: SingleType a -> ((a, a) -> a)
evalMin :: forall a. SingleType a -> (a, a) -> a
evalMin (NumSingleType (IntegralNumType IntegralType a
ty)) | IntegralDict a
IntegralDict <- IntegralType a -> IntegralDict a
forall a. IntegralType a -> IntegralDict a
integralDict IntegralType a
ty = (a -> a -> a) -> (a, a) -> a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> a
forall a. Ord a => a -> a -> a
min
evalMin (NumSingleType (FloatingNumType FloatingType a
ty)) | FloatingDict a
FloatingDict <- FloatingType a -> FloatingDict a
forall a. FloatingType a -> FloatingDict a
floatingDict FloatingType a
ty = (a -> a -> a) -> (a, a) -> a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> a
forall a. Ord a => a -> a -> a
min


{--
-- Sequence evaluation
-- ---------------

-- Position in sequence.
--
type SeqPos = Int

-- Configuration for sequence evaluation.
--
data SeqConfig = SeqConfig
  { chunkSize :: Int -- Allocation limit for a sequence in
                     -- words. Actual runtime allocation should be the
                     -- maximum of this size and the size of the
                     -- largest element in the sequence.
  }

-- Default sequence evaluation configuration for testing purposes.
--
defaultSeqConfig :: SeqConfig
defaultSeqConfig = SeqConfig { chunkSize = 2 }

type Chunk a = Vector' a

-- The empty chunk. O(1).
emptyChunk :: Arrays a => Chunk a
emptyChunk = empty'

-- Number of arrays in chunk. O(1).
--
clen :: Arrays a => Chunk a -> Int
clen = length'

elemsPerChunk :: SeqConfig -> Int -> Int
elemsPerChunk conf n
  | n < 1 = chunkSize conf
  | otherwise =
    let (a,b) = chunkSize conf `quotRem` n
    in a + signum b

-- Drop a number of arrays from a chunk. O(1). Note: Require keeping a
-- scan of element sizes.
--
cdrop :: Arrays a => Int -> Chunk a -> Chunk a
cdrop = drop' dropOp (fst . offsetsOp)

-- Get all the shapes of a chunk of arrays. O(1).
--
chunkShapes :: Chunk (Array sh a) -> Vector sh
chunkShapes = shapes'

-- Get all the elements of a chunk of arrays. O(1).
--
chunkElems :: Chunk (Array sh a) -> Vector a
chunkElems = elements'

-- Convert a vector to a chunk of scalars.
--
vec2Chunk :: Elt e => Vector e -> Chunk (Scalar e)
vec2Chunk = vec2Vec'

-- Convert a list of arrays to a chunk.
--
fromListChunk :: Arrays a => [a] -> Vector' a
fromListChunk = fromList' concatOp

-- Convert a chunk to a list of arrays.
--
toListChunk :: Arrays a => Vector' a -> [a]
toListChunk = toList' fetchAllOp

-- fmap for Chunk. O(n).
--   TODO: Use vectorised function.
mapChunk :: (Arrays a, Arrays b)
         => (a -> b)
         -> Chunk a -> Chunk b
mapChunk f c = fromListChunk $ map f (toListChunk c)

-- zipWith for Chunk. O(n).
--  TODO: Use vectorised function.
zipWithChunk :: (Arrays a, Arrays b, Arrays c)
             => (a -> b -> c)
             -> Chunk a -> Chunk b -> Chunk c
zipWithChunk f c1 c2 = fromListChunk $ zipWith f (toListChunk c1) (toListChunk c2)

-- A window on a sequence.
--
data Window a = Window
  { chunk :: Chunk a   -- Current allocated chunk.
  , wpos  :: SeqPos    -- Position of the window on the sequence, given
                       -- in number of elements.
  }

-- The initial empty window.
--
window0 :: Arrays a => Window a
window0 = Window { chunk = emptyChunk, wpos = 0 }

-- Index the given window by the given index on the sequence.
--
(!#) :: Arrays a => Window a -> SeqPos -> Chunk a
w !# i
  | j <- i - wpos w
  , j >= 0
  = cdrop j (chunk w)
  --
  | otherwise
  = error $ "Window indexed before position. wpos = " ++ show (wpos w) ++ " i = " ++ show i

-- Move the give window by supplying the next chunk.
--
moveWin :: Arrays a => Window a -> Chunk a -> Window a
moveWin w c = w { chunk = c
                , wpos = wpos w + clen (chunk w)
                }

-- A cursor on a sequence.
--
data Cursor senv a = Cursor
  { ref  :: Idx senv a -- Reference to the sequence.
  , cpos :: SeqPos     -- Position of the cursor on the sequence,
                       -- given in number of elements.
  }

-- Initial cursor.
--
cursor0 :: Idx senv a -> Cursor senv a
cursor0 x = Cursor { ref = x, cpos = 0 }

-- Advance cursor by a relative amount.
--
moveCursor :: Int -> Cursor senv a -> Cursor senv a
moveCursor k c = c { cpos = cpos c + k }

-- Valuation for an environment of sequence windows.
--
data Val' senv where
  Empty' :: Val' ()
  Push'  :: Val' senv -> Window t -> Val' (senv, t)

-- Projection of a window from a window valuation using a de Bruijn
-- index.
--
prj' :: Idx senv t -> Val' senv -> Window t
prj' ZeroIdx       (Push' _   v) = v
prj' (SuccIdx idx) (Push' val _) = prj' idx val

-- Projection of a chunk from a window valuation using a sequence
-- cursor.
--
prjChunk :: Arrays a => Cursor senv a -> Val' senv -> Chunk a
prjChunk c senv = prj' (ref c) senv !# cpos c

-- An executable sequence.
--
data ExecSeq senv arrs where
  ExecP :: Arrays a => Window a -> ExecP senv a -> ExecSeq (senv, a) arrs -> ExecSeq senv  arrs
  ExecC :: Arrays a =>             ExecC senv a ->                           ExecSeq senv  a
  ExecR :: Arrays a =>             Cursor senv a ->                          ExecSeq senv  [a]

-- An executable producer.
--
data ExecP senv a where
  ExecStreamIn :: Int
               -> [a]
               -> ExecP senv a

  ExecMap :: Arrays a
          => (Chunk a -> Chunk b)
          -> Cursor senv a
          -> ExecP senv b

  ExecZipWith :: (Arrays a, Arrays b)
              => (Chunk a -> Chunk b -> Chunk c)
              -> Cursor senv a
              -> Cursor senv b
              -> ExecP senv c

  -- Stream scan skeleton.
  ExecScan :: Arrays a
           => (s -> Chunk a -> (Chunk r, s)) -- Chunk scanner.
           -> s                              -- Accumulator (internal state).
           -> Cursor senv a                  -- Input stream.
           -> ExecP senv r

-- An executable consumer.
--
data ExecC senv a where

  -- Stream reduction skeleton.
  ExecFold :: Arrays a
           => (s -> Chunk a -> s) -- Chunk consumer function.
           -> (s -> r)            -- Finalizer function.
           -> s                   -- Accumulator (internal state).
           -> Cursor senv a       -- Input stream.
           -> ExecC senv r

  ExecStuple :: IsAtuple a
             => Atuple (ExecC senv) (TupleRepr a)
             -> ExecC senv a

minCursor :: ExecSeq senv a -> SeqPos
minCursor s = travS s 0
  where
    travS :: ExecSeq senv a -> Int -> SeqPos
    travS s i =
      case s of
        ExecP _ p s' -> travP p i `min` travS s' (i+1)
        ExecC   c    -> travC c i
        ExecR   _    -> maxBound

    k :: Cursor senv a -> Int -> SeqPos
    k c i
      | i == idxToInt (ref c) = cpos c
      | otherwise             = maxBound

    travP :: ExecP senv a -> Int -> SeqPos
    travP p i =
      case p of
        ExecStreamIn _ _ -> maxBound
        ExecMap _ c -> k c i
        ExecZipWith _ c1 c2 -> k c1 i `min` k c2 i
        ExecScan _ _ c -> k c i

    travT :: Atuple (ExecC senv) t -> Int -> SeqPos
    travT NilAtup        _ = maxBound
    travT (SnocAtup t c) i = travT t i `min` travC c i

    travC :: ExecC senv a -> Int -> SeqPos
    travC c i =
      case c of
        ExecFold _ _ _ cu -> k cu i
        ExecStuple t      -> travT t i


evalDelayedSeq
    :: SeqConfig
    -> DelayedSeq arrs
    -> arrs
evalDelayedSeq cfg (DelayedSeq aenv s) | aenv' <- evalExtend aenv Empty
                                       = evalSeq cfg s aenv'

evalSeq :: forall aenv arrs.
            SeqConfig
         -> PreOpenSeq DelayedOpenAcc aenv () arrs
         -> Val aenv -> arrs
evalSeq conf s aenv = evalSeq' s
  where
    evalSeq' :: PreOpenSeq DelayedOpenAcc aenv senv arrs -> arrs
    evalSeq' (Producer _ s) = evalSeq' s
    evalSeq' (Consumer _)   = loop (initSeq aenv s)
    evalSeq' (Reify _)      = reify (initSeq aenv s)

    -- Initialize the producers and the accumulators of the consumers
    -- with the given array enviroment.
    initSeq :: forall senv arrs'.
                Val aenv
             -> PreOpenSeq DelayedOpenAcc aenv senv arrs'
             -> ExecSeq senv arrs'
    initSeq aenv s =
      case s of
        Producer   p s' -> ExecP window0 (initProducer p) (initSeq aenv s')
        Consumer   c    -> ExecC         (initConsumer c)
        Reify      ix   -> ExecR (cursor0 ix)

    -- Generate a list from the sequence.
    reify :: forall arrs. ExecSeq () [arrs]
          -> [arrs]
    reify s = case step s Empty' of
                (Just s', a) -> a ++ reify s'
                (Nothing, a) -> a

    -- Iterate the given sequence until it terminates.
    -- A sequence only terminates when one of the producers are exhausted.
    loop :: Arrays arrs
         => ExecSeq () arrs
         -> arrs
    loop s =
      case step' s of
        (Nothing, arrs) -> arrs
        (Just s', _)    -> loop s'

      where
        step' :: ExecSeq () arrs -> (Maybe (ExecSeq () arrs), arrs)
        step' s = step s Empty'

    -- One iteration of a sequence.
    step :: forall senv arrs'.
            ExecSeq senv arrs'
         -> Val' senv
         -> (Maybe (ExecSeq senv arrs'), arrs')
    step s senv =
      case s of
        ExecP w p s' ->
          let (c, mp')  = produce p senv
              finished  = 0 == clen (w !# minCursor s')
              w'        = if finished then moveWin w c else w
              (ms'', a) = step s' (senv `Push'` w')
          in case ms'' of
            Nothing  -> (Nothing, a)
            Just s'' | finished
                     , Just p' <- mp'
                     -> (Just (ExecP w' p' s''), a)
                     | not finished
                     -> (Just (ExecP w' p  s''), a)
                     | otherwise
                     -> (Nothing, a)
        ExecC   c    -> let (c', acc) = consume c senv
                        in (Just (ExecC c'), acc)
        ExecR ix     -> let c = prjChunk ix senv in (Just (ExecR (moveCursor (clen c) ix)), toListChunk c)

    evalA :: DelayedOpenAcc aenv a -> a
    evalA acc = evalOpenAcc acc aenv

    evalAF :: DelayedOpenAfun aenv f -> f
    evalAF f = evalOpenAfun f aenv

    evalE :: DelayedExp aenv t -> t
    evalE exp = evalExp exp aenv

    evalF :: DelayedFun aenv f -> f
    evalF fun = evalFun fun aenv

    initProducer :: forall a senv.
                    Producer DelayedOpenAcc aenv senv a
                 -> ExecP senv a
    initProducer p =
      case p of
        StreamIn arrs -> ExecStreamIn 1 arrs
        ToSeq sliceIndex slix (delayed -> Delayed sh ix _) ->
          let n   = R.size (R.sliceShape sliceIndex (fromElt sh))
              k   = elemsPerChunk conf n
          in ExecStreamIn k (toSeqOp sliceIndex slix (fromFunction sh ix))
        MapSeq     f x       -> ExecMap     (mapChunk (evalAF f)) (cursor0 x)
        ChunkedMapSeq f x    -> ExecMap     (evalAF f) (cursor0 x)
        ZipWithSeq f x y     -> ExecZipWith (zipWithChunk (evalAF f)) (cursor0 x) (cursor0 y)
        ScanSeq    f e x     -> ExecScan scanner (evalE e) (cursor0 x)
          where
            scanner a c =
              let v0 = chunkElems c
                  (v1, a') = scanl'Op (evalF f) a (delayArray v0)
              in (vec2Chunk v1, fromScalar a')

    initConsumer :: forall a senv.
                    Consumer DelayedOpenAcc aenv senv a
                 -> ExecC senv a
    initConsumer c =
      case c of
        FoldSeq f e x ->
          let f' = evalF f
              a0 = fromFunction (Z :. chunkSize conf) (const (evalE e))
              consumer v c = zipWith'Op f' (delayArray v) (delayArray (chunkElems c))
              finalizer = fold1Op f' . delayArray
          in ExecFold consumer finalizer a0 (cursor0 x)
        FoldSeqFlatten f acc x ->
          let f' = evalAF f
              a0 = evalA acc
              consumer a c = f' a (chunkShapes c) (chunkElems c)
          in ExecFold consumer id a0 (cursor0 x)
        Stuple t ->
          let initTup :: Atuple (Consumer DelayedOpenAcc aenv senv) t -> Atuple (ExecC senv) t
              initTup NilAtup        = NilAtup
              initTup (SnocAtup t c) = SnocAtup (initTup t) (initConsumer c)
          in ExecStuple (initTup t)

    delayed :: DelayedOpenAcc aenv (Array sh e) -> Delayed (Array sh e)
    delayed AST.Manifest{}  = $internalError "evalOpenAcc" "expected delayed array"
    delayed AST.Delayed{..} = Delayed (evalExp extentD aenv)
                                      (evalFun indexD aenv)
                                      (evalFun linearIndexD aenv)

produce :: Arrays a => ExecP senv a -> Val' senv -> (Chunk a, Maybe (ExecP senv a))
produce p senv =
  case p of
    ExecStreamIn k xs ->
      let (xs', xs'') = (take k xs, drop k xs)
          c           = fromListChunk xs'
          mp          = if null xs''
                        then Nothing
                        else Just (ExecStreamIn k xs'')
      in (c, mp)
    ExecMap f x ->
      let c = prjChunk x senv
      in (f c, Just $ ExecMap f (moveCursor (clen c) x))
    ExecZipWith f x y ->
      let c1 = prjChunk x senv
          c2 = prjChunk y senv
          k = clen c1 `min` clen c2
      in (f c1 c2, Just $ ExecZipWith f (moveCursor k x) (moveCursor k y))
    ExecScan scanner a x ->
      let c = prjChunk x senv
          (c', a') = scanner a c
          k = clen c
      in (c', Just $ ExecScan scanner a' (moveCursor k x))

consume :: forall senv a. ExecC senv a -> Val' senv -> (ExecC senv a, a)
consume c senv =
  case c of
    ExecFold f g acc x ->
      let c    = prjChunk x senv
          acc' = f acc c
      -- Even though we call g here, lazy evaluation should guarantee it is
      -- only ever called once.
      in (ExecFold f g acc' (moveCursor (clen c) x), g acc')
    ExecStuple t ->
      let consT :: Atuple (ExecC senv) t -> (Atuple (ExecC senv) t, t)
          consT NilAtup        = (NilAtup, ())
          consT (SnocAtup t c) | (c', acc) <- consume c senv
                               , (t', acc') <- consT t
                               = (SnocAtup t' c', (acc', acc))
          (t', acc) = consT t
      in (ExecStuple t', toAtuple acc)

evalExtend :: Extend DelayedOpenAcc aenv aenv' -> Val aenv -> Val aenv'
evalExtend BaseEnv aenv = aenv
evalExtend (PushEnv ext1 ext2) aenv | aenv' <- evalExtend ext1 aenv
                                    = Push aenv' (evalOpenAcc ext2 aenv')

delayArray :: Array sh e -> Delayed (Array sh e)
delayArray arr@(Array _ adata) = Delayed (shape arr) (arr!) (toElt . unsafeIndexArrayData adata)

fromScalar :: Scalar a -> a
fromScalar = (!Z)

concatOp :: forall e. Elt e => [Vector e] -> Vector e
concatOp = concatVectors

fetchAllOp :: (Shape sh, Elt e) => Segments sh -> Vector e -> [Array sh e]
fetchAllOp segs elts
  | (offsets, n) <- offsetsOp segs
  , (n ! Z) <= size (shape elts)
  = [fetch (segs ! (Z :. i)) (offsets ! (Z :. i)) | i <- [0 .. size (shape segs) - 1]]
  | otherwise = error $ "illegal argument to fetchAllOp"
  where
    fetch sh offset = fromFunction sh (\ ix -> elts ! (Z :. ((toIndex sh ix) + offset)))

dropOp :: Elt e => Int -> Vector e -> Vector e
dropOp i v   -- TODO
             --  * Implement using C-style pointer-plus.
             --    ; dropOp is used often (from prjChunk),
             --      so it ought to be efficient O(1).
  | n <- size (shape v)
  , i <= n
  , i >= 0
  = fromFunction (Z :. n - i) (\ (Z :. j) -> v ! (Z :. i + j))
  | otherwise = error $ "illegal argument to drop"

offsetsOp :: Shape sh => Segments sh -> (Vector Int, Scalar Int)
offsetsOp segs = scanl'Op (+) 0 $ delayArray (mapOp size (delayArray segs))
--}