{-# LANGUAGE Rank2Types #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
module Data.Array.Knead.Symbolic.Private where

import qualified Data.Array.Knead.Shape as Shape
import qualified Data.Array.Knead.Expression as Expr

import LLVM.DSL.Expression (Exp(Exp))

import qualified LLVM.Extra.Multi.Value as MultiValue
import qualified LLVM.Extra.Iterator as Iter
import qualified LLVM.Extra.Maybe as Maybe
import qualified LLVM.Core as LLVM

import qualified Control.Category as Cat
import qualified Control.Monad.HT as Monad
import Control.Monad ((<=<), )

import Prelude hiding (id, map, zipWith, replicate, )


type Val = MultiValue.T
type Code r a = LLVM.CodeGenFunction r (Val a)

data Array sh a =
   Array (Exp sh) (forall r. Val (Shape.Index sh) -> Code r a)

shape :: Array sh a -> Exp sh
shape :: forall sh a. Array sh a -> Exp sh
shape (Array Exp sh
sh forall r. Val (Index sh) -> Code r a
_) = Exp sh
sh

(!) ::
   (Shape.C sh,  Shape.Index sh  ~ ix) =>
   Array sh a -> Exp ix -> Exp a
! :: forall sh ix a.
(C sh, Index sh ~ ix) =>
Array sh a -> Exp ix -> Exp a
(!) (Array Exp sh
_ forall r. Val (Index sh) -> Code r a
code) (Exp forall r. CodeGenFunction r (T ix)
ix) = (forall r. CodeGenFunction r (T a)) -> Exp a
forall a. (forall r. CodeGenFunction r (T a)) -> Exp a
Exp (T ix -> CodeGenFunction r (T a)
Val (Index sh) -> CodeGenFunction r (T a)
forall r. Val (Index sh) -> Code r a
code (T ix -> CodeGenFunction r (T a))
-> CodeGenFunction r (T ix) -> CodeGenFunction r (T a)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< CodeGenFunction r (T ix)
forall r. CodeGenFunction r (T ix)
ix)

the :: (Shape.Scalar sh) => Array sh a -> Exp a
the :: forall sh a. Scalar sh => Array sh a -> Exp a
the (Array Exp sh
z forall r. Val (Index sh) -> Code r a
code) = (forall r. CodeGenFunction r (T a)) -> Exp a
forall a. (forall r. CodeGenFunction r (T a)) -> Exp a
Exp (Val (Index sh) -> Code r a
forall r. Val (Index sh) -> Code r a
code (Val (Index sh) -> Code r a) -> Val (Index sh) -> Code r a
forall a b. (a -> b) -> a -> b
$ Exp sh -> Val (Index sh)
forall sh (val :: * -> *) (f :: * -> *).
(Scalar sh, Value val) =>
f sh -> val (Index sh)
forall (val :: * -> *) (f :: * -> *).
Value val =>
f sh -> val (Index sh)
Shape.zeroIndex Exp sh
z)

fromScalar :: (Shape.Scalar sh) => Exp a -> Array sh a
fromScalar :: forall sh a. Scalar sh => Exp a -> Array sh a
fromScalar = Exp sh -> Exp a -> Array sh a
forall sh a. Exp sh -> Exp a -> Array sh a
fill Exp sh
forall sh (val :: * -> *). (Scalar sh, Value val) => val sh
forall (val :: * -> *). Value val => val sh
Shape.scalar


fill :: Exp sh -> Exp a -> Array sh a
fill :: forall sh a. Exp sh -> Exp a -> Array sh a
fill Exp sh
sh (Exp forall r. CodeGenFunction r (T a)
code) = Exp sh -> (forall r. Val (Index sh) -> Code r a) -> Array sh a
forall sh a.
Exp sh -> (forall r. Val (Index sh) -> Code r a) -> Array sh a
Array Exp sh
sh (\Val (Index sh)
_z -> Code r a
forall r. CodeGenFunction r (T a)
code)


{- |
This class allows to implement functions without parameters
for both simple and parameterized arrays.
-}
class C array where
   lift0 :: Array sh a -> array sh a
   lift1 :: (Array sha a -> Array shb b) -> array sha a -> array shb b
   lift2 ::
      (Array sha a -> Array shb b -> Array shc c) ->
      array sha a -> array shb b -> array shc c

instance C Array where
   lift0 :: forall sh a. Array sh a -> Array sh a
lift0 = Array sh a -> Array sh a
forall a. a -> a
forall {k} (cat :: k -> k -> *) (a :: k). Category cat => cat a a
Cat.id
   lift1 :: forall sha a shb b.
(Array sha a -> Array shb b) -> Array sha a -> Array shb b
lift1 = (Array sha a -> Array shb b) -> Array sha a -> Array shb b
forall a. a -> a
forall {k} (cat :: k -> k -> *) (a :: k). Category cat => cat a a
Cat.id
   lift2 :: forall sha a shb b shc c.
(Array sha a -> Array shb b -> Array shc c)
-> Array sha a -> Array shb b -> Array shc c
lift2 = (Array sha a -> Array shb b -> Array shc c)
-> Array sha a -> Array shb b -> Array shc c
forall a. a -> a
forall {k} (cat :: k -> k -> *) (a :: k). Category cat => cat a a
Cat.id


gather ::
   (C array,
    Shape.C sh0, Shape.Index sh0 ~ ix0,
    Shape.C sh1, Shape.Index sh1 ~ ix1,
    MultiValue.C a) =>
   array sh1 ix0 ->
   array sh0 a ->
   array sh1 a
gather :: forall (array :: * -> * -> *) sh0 ix0 sh1 ix1 a.
(C array, C sh0, Index sh0 ~ ix0, C sh1, Index sh1 ~ ix1, C a) =>
array sh1 ix0 -> array sh0 a -> array sh1 a
gather =
   (Array sh1 ix0 -> Array sh0 a -> Array sh1 a)
-> array sh1 ix0 -> array sh0 a -> array sh1 a
forall sha a shb b shc c.
(Array sha a -> Array shb b -> Array shc c)
-> array sha a -> array shb b -> array shc c
forall (array :: * -> * -> *) sha a shb b shc c.
C array =>
(Array sha a -> Array shb b -> Array shc c)
-> array sha a -> array shb b -> array shc c
lift2 ((Array sh1 ix0 -> Array sh0 a -> Array sh1 a)
 -> array sh1 ix0 -> array sh0 a -> array sh1 a)
-> (Array sh1 ix0 -> Array sh0 a -> Array sh1 a)
-> array sh1 ix0
-> array sh0 a
-> array sh1 a
forall a b. (a -> b) -> a -> b
$ \(Array Exp sh1
sh1 forall r. Val (Index sh1) -> Code r ix0
f) (Array Exp sh0
_sh0 forall r. Val (Index sh0) -> Code r a
code) ->
      Exp sh1 -> (forall r. Val (Index sh1) -> Code r a) -> Array sh1 a
forall sh a.
Exp sh -> (forall r. Val (Index sh) -> Code r a) -> Array sh a
Array Exp sh1
sh1 (T ix0 -> CodeGenFunction r (Val a)
Val (Index sh0) -> CodeGenFunction r (Val a)
forall r. Val (Index sh0) -> Code r a
code (T ix0 -> CodeGenFunction r (Val a))
-> (Val ix1 -> CodeGenFunction r (T ix0))
-> Val ix1
-> CodeGenFunction r (Val a)
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< Val ix1 -> CodeGenFunction r (T ix0)
Val (Index sh1) -> CodeGenFunction r (T ix0)
forall r. Val (Index sh1) -> Code r ix0
f)

backpermute2 ::
   (C array,
    Shape.C sh0, Shape.Index sh0 ~ ix0,
    Shape.C sh1, Shape.Index sh1 ~ ix1,
    Shape.C sh,  Shape.Index sh  ~ ix) =>
   Exp sh ->
   (Exp ix -> Exp ix0) ->
   (Exp ix -> Exp ix1) ->
   (Exp a -> Exp b -> Exp c) ->
   array sh0 a -> array sh1 b -> array sh c
backpermute2 :: forall (array :: * -> * -> *) sh0 ix0 sh1 ix1 sh ix a b c.
(C array, C sh0, Index sh0 ~ ix0, C sh1, Index sh1 ~ ix1, C sh,
 Index sh ~ ix) =>
Exp sh
-> (Exp ix -> Exp ix0)
-> (Exp ix -> Exp ix1)
-> (Exp a -> Exp b -> Exp c)
-> array sh0 a
-> array sh1 b
-> array sh c
backpermute2 Exp sh
sh Exp ix -> Exp ix0
projectIndex0 Exp ix -> Exp ix1
projectIndex1 Exp a -> Exp b -> Exp c
f =
   (Array sh0 a -> Array sh1 b -> Array sh c)
-> array sh0 a -> array sh1 b -> array sh c
forall sha a shb b shc c.
(Array sha a -> Array shb b -> Array shc c)
-> array sha a -> array shb b -> array shc c
forall (array :: * -> * -> *) sha a shb b shc c.
C array =>
(Array sha a -> Array shb b -> Array shc c)
-> array sha a -> array shb b -> array shc c
lift2 ((Array sh0 a -> Array sh1 b -> Array sh c)
 -> array sh0 a -> array sh1 b -> array sh c)
-> (Array sh0 a -> Array sh1 b -> Array sh c)
-> array sh0 a
-> array sh1 b
-> array sh c
forall a b. (a -> b) -> a -> b
$ \(Array Exp sh0
_sha forall r. Val (Index sh0) -> Code r a
codeA) (Array Exp sh1
_shb forall r. Val (Index sh1) -> Code r b
codeB) ->
      Exp sh -> (forall r. Val (Index sh) -> Code r c) -> Array sh c
forall sh a.
Exp sh -> (forall r. Val (Index sh) -> Code r a) -> Array sh a
Array Exp sh
sh
         (\Val (Index sh)
ix ->
            (Val a -> Val b -> Code r c)
-> CodeGenFunction r (Val a)
-> CodeGenFunction r (Val b)
-> Code r c
forall (m :: * -> *) a b c.
Monad m =>
(a -> b -> m c) -> m a -> m b -> m c
Monad.liftJoin2 ((Exp a -> Exp b -> Exp c) -> Val a -> Val b -> Code r c
forall ae am be bm ce cm r.
(Aggregate ae am, Aggregate be bm, Aggregate ce cm) =>
(ae -> be -> ce) -> am -> bm -> CodeGenFunction r cm
Expr.unliftM2 Exp a -> Exp b -> Exp c
f)
               (Val ix0 -> CodeGenFunction r (Val a)
Val (Index sh0) -> CodeGenFunction r (Val a)
forall r. Val (Index sh0) -> Code r a
codeA (Val ix0 -> CodeGenFunction r (Val a))
-> CodeGenFunction r (Val ix0) -> CodeGenFunction r (Val a)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Exp ix -> Exp ix0) -> Val ix -> CodeGenFunction r (Val ix0)
forall ae am be bm r.
(Aggregate ae am, Aggregate be bm) =>
(ae -> be) -> am -> CodeGenFunction r bm
Expr.unliftM1 Exp ix -> Exp ix0
projectIndex0 Val ix
Val (Index sh)
ix)
               (Val ix1 -> CodeGenFunction r (Val b)
Val (Index sh1) -> CodeGenFunction r (Val b)
forall r. Val (Index sh1) -> Code r b
codeB (Val ix1 -> CodeGenFunction r (Val b))
-> CodeGenFunction r (Val ix1) -> CodeGenFunction r (Val b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (Exp ix -> Exp ix1) -> Val ix -> CodeGenFunction r (Val ix1)
forall ae am be bm r.
(Aggregate ae am, Aggregate be bm) =>
(ae -> be) -> am -> CodeGenFunction r bm
Expr.unliftM1 Exp ix -> Exp ix1
projectIndex1 Val ix
Val (Index sh)
ix))


id ::
   (C array, Shape.C sh, Shape.Index sh ~ ix) =>
   Exp sh -> array sh ix
id :: forall (array :: * -> * -> *) sh ix.
(C array, C sh, Index sh ~ ix) =>
Exp sh -> array sh ix
id Exp sh
sh = Array sh ix -> array sh ix
forall sh a. Array sh a -> array sh a
forall (array :: * -> * -> *) sh a.
C array =>
Array sh a -> array sh a
lift0 (Array sh ix -> array sh ix) -> Array sh ix -> array sh ix
forall a b. (a -> b) -> a -> b
$ Exp sh -> (forall r. Val (Index sh) -> Code r ix) -> Array sh ix
forall sh a.
Exp sh -> (forall r. Val (Index sh) -> Code r a) -> Array sh a
Array Exp sh
sh Val ix -> CodeGenFunction r (Val ix)
Val (Index sh) -> CodeGenFunction r (Val ix)
forall a. a -> CodeGenFunction r a
forall r. Val (Index sh) -> Code r ix
forall (m :: * -> *) a. Monad m => a -> m a
return

map ::
   (C array, Shape.C sh) =>
   (Exp a -> Exp b) ->
   array sh a -> array sh b
map :: forall (array :: * -> * -> *) sh a b.
(C array, C sh) =>
(Exp a -> Exp b) -> array sh a -> array sh b
map Exp a -> Exp b
f =
   (Array sh a -> Array sh b) -> array sh a -> array sh b
forall sha a shb b.
(Array sha a -> Array shb b) -> array sha a -> array shb b
forall (array :: * -> * -> *) sha a shb b.
C array =>
(Array sha a -> Array shb b) -> array sha a -> array shb b
lift1 ((Array sh a -> Array sh b) -> array sh a -> array sh b)
-> (Array sh a -> Array sh b) -> array sh a -> array sh b
forall a b. (a -> b) -> a -> b
$ \(Array Exp sh
sh forall r. Val (Index sh) -> Code r a
code) ->
      Exp sh -> (forall r. Val (Index sh) -> Code r b) -> Array sh b
forall sh a.
Exp sh -> (forall r. Val (Index sh) -> Code r a) -> Array sh a
Array Exp sh
sh ((Exp a -> Exp b) -> Val a -> CodeGenFunction r (Val b)
forall ae am be bm r.
(Aggregate ae am, Aggregate be bm) =>
(ae -> be) -> am -> CodeGenFunction r bm
Expr.unliftM1 Exp a -> Exp b
f (Val a -> CodeGenFunction r (Val b))
-> (Val (Index sh) -> CodeGenFunction r (Val a))
-> Val (Index sh)
-> CodeGenFunction r (Val b)
forall (m :: * -> *) b c a.
Monad m =>
(b -> m c) -> (a -> m b) -> a -> m c
<=< Val (Index sh) -> CodeGenFunction r (Val a)
forall r. Val (Index sh) -> Code r a
code)

mapWithIndex ::
   (C array, Shape.C sh, Shape.Index sh ~ ix) =>
   (Exp ix -> Exp a -> Exp b) ->
   array sh a -> array sh b
mapWithIndex :: forall (array :: * -> * -> *) sh ix a b.
(C array, C sh, Index sh ~ ix) =>
(Exp ix -> Exp a -> Exp b) -> array sh a -> array sh b
mapWithIndex Exp ix -> Exp a -> Exp b
f =
   (Array sh a -> Array sh b) -> array sh a -> array sh b
forall sha a shb b.
(Array sha a -> Array shb b) -> array sha a -> array shb b
forall (array :: * -> * -> *) sha a shb b.
C array =>
(Array sha a -> Array shb b) -> array sha a -> array shb b
lift1 ((Array sh a -> Array sh b) -> array sh a -> array sh b)
-> (Array sh a -> Array sh b) -> array sh a -> array sh b
forall a b. (a -> b) -> a -> b
$ \(Array Exp sh
sh forall r. Val (Index sh) -> Code r a
code) ->
      Exp sh -> (forall r. Val (Index sh) -> Code r b) -> Array sh b
forall sh a.
Exp sh -> (forall r. Val (Index sh) -> Code r a) -> Array sh a
Array Exp sh
sh (\Val (Index sh)
ix -> (Exp ix -> Exp a -> Exp b) -> Val ix -> Val a -> Code r b
forall ae am be bm ce cm r.
(Aggregate ae am, Aggregate be bm, Aggregate ce cm) =>
(ae -> be -> ce) -> am -> bm -> CodeGenFunction r cm
Expr.unliftM2 Exp ix -> Exp a -> Exp b
f Val ix
Val (Index sh)
ix (Val a -> Code r b) -> CodeGenFunction r (Val a) -> Code r b
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Val (Index sh) -> CodeGenFunction r (Val a)
forall r. Val (Index sh) -> Code r a
code Val (Index sh)
ix)


fold1Code ::
   (Shape.C sh, Shape.Index sh ~ ix, MultiValue.C a) =>
   (Exp a -> Exp a -> Exp a) ->
   Exp sh ->
   (Val ix -> Code r a) ->
   Code r a
fold1Code :: forall sh ix a r.
(C sh, Index sh ~ ix, C a) =>
(Exp a -> Exp a -> Exp a)
-> Exp sh -> (Val ix -> Code r a) -> Code r a
fold1Code Exp a -> Exp a -> Exp a
f (Exp forall r. CodeGenFunction r (T sh)
nc) Val ix -> Code r a
code = do
   T sh
n <- CodeGenFunction r (T sh)
forall r. CodeGenFunction r (T sh)
nc
   (T (Val a) -> Val a) -> CodeGenFunction r (T (Val a)) -> Code r a
forall a b. (a -> b) -> CodeGenFunction r a -> CodeGenFunction r b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap T (Val a) -> Val a
forall a. T a -> a
Maybe.fromJust (CodeGenFunction r (T (Val a)) -> Code r a)
-> CodeGenFunction r (T (Val a)) -> Code r a
forall a b. (a -> b) -> a -> b
$
      (Val ix -> T (Val a) -> CodeGenFunction r (T (Val a)))
-> T sh -> T (Val a) -> CodeGenFunction r (T (Val a))
forall sh ix state r.
(C sh, Index sh ~ ix, Phi state) =>
(T ix -> state -> CodeGenFunction r state)
-> T sh -> state -> CodeGenFunction r state
forall ix state r.
(Index sh ~ ix, Phi state) =>
(T ix -> state -> CodeGenFunction r state)
-> T sh -> state -> CodeGenFunction r state
Shape.loop
         (\Val ix
i0 T (Val a)
macc0 -> do
            Val a
a <- Val ix -> Code r a
code Val ix
i0
            Val a
acc1 <- T (Val a) -> Code r a -> (Val a -> Code r a) -> Code r a
forall b a r.
Phi b =>
T a
-> CodeGenFunction r b
-> (a -> CodeGenFunction r b)
-> CodeGenFunction r b
Maybe.run T (Val a)
macc0 (Val a -> Code r a
forall a. a -> CodeGenFunction r a
forall (m :: * -> *) a. Monad m => a -> m a
return Val a
a) ((Val a -> Val a -> Code r a) -> Val a -> Val a -> Code r a
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Exp a -> Exp a -> Exp a) -> Val a -> Val a -> Code r a
forall ae am be bm ce cm r.
(Aggregate ae am, Aggregate be bm, Aggregate ce cm) =>
(ae -> be -> ce) -> am -> bm -> CodeGenFunction r cm
Expr.unliftM2 Exp a -> Exp a -> Exp a
f) Val a
a)
            T (Val a) -> CodeGenFunction r (T (Val a))
forall a. a -> CodeGenFunction r a
forall (m :: * -> *) a. Monad m => a -> m a
return (T (Val a) -> CodeGenFunction r (T (Val a)))
-> T (Val a) -> CodeGenFunction r (T (Val a))
forall a b. (a -> b) -> a -> b
$ Val a -> T (Val a)
forall a. a -> T a
Maybe.just Val a
acc1)
         T sh
n T (Val a)
forall a. Undefined a => T a
Maybe.nothing

fold1 ::
   (C array, Shape.C sh0, Shape.C sh1, MultiValue.C a) =>
   (Exp a -> Exp a -> Exp a) ->
   array (sh0, sh1) a -> array sh0 a
fold1 :: forall (array :: * -> * -> *) sh0 sh1 a.
(C array, C sh0, C sh1, C a) =>
(Exp a -> Exp a -> Exp a) -> array (sh0, sh1) a -> array sh0 a
fold1 Exp a -> Exp a -> Exp a
f =
   (Array (sh0, sh1) a -> Array sh0 a)
-> array (sh0, sh1) a -> array sh0 a
forall sha a shb b.
(Array sha a -> Array shb b) -> array sha a -> array shb b
forall (array :: * -> * -> *) sha a shb b.
C array =>
(Array sha a -> Array shb b) -> array sha a -> array shb b
lift1 ((Array (sh0, sh1) a -> Array sh0 a)
 -> array (sh0, sh1) a -> array sh0 a)
-> (Array (sh0, sh1) a -> Array sh0 a)
-> array (sh0, sh1) a
-> array sh0 a
forall a b. (a -> b) -> a -> b
$ \(Array Exp (sh0, sh1)
shs forall r. Val (Index (sh0, sh1)) -> Code r a
code) ->
      case Exp (sh0, sh1) -> (Exp sh0, Exp sh1)
forall (val :: * -> *) a b.
Value val =>
val (a, b) -> (val a, val b)
Expr.unzip Exp (sh0, sh1)
shs of
         (Exp sh0
sh, Exp sh1
s) -> Exp sh0 -> (forall {r}. Val (Index sh0) -> Code r a) -> Array sh0 a
forall sh a.
Exp sh -> (forall r. Val (Index sh) -> Code r a) -> Array sh a
Array Exp sh0
sh ((forall {r}. Val (Index sh0) -> Code r a) -> Array sh0 a)
-> (forall {r}. Val (Index sh0) -> Code r a) -> Array sh0 a
forall a b. (a -> b) -> a -> b
$ (Exp a -> Exp a -> Exp a)
-> Exp sh1 -> (Val (Index sh1) -> Code r a) -> Code r a
forall sh ix a r.
(C sh, Index sh ~ ix, C a) =>
(Exp a -> Exp a -> Exp a)
-> Exp sh -> (Val ix -> Code r a) -> Code r a
fold1Code Exp a -> Exp a -> Exp a
f Exp sh1
s ((Val (Index sh1) -> Code r a) -> Code r a)
-> (Val (Index sh0) -> Val (Index sh1) -> Code r a)
-> Val (Index sh0)
-> Code r a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (T (Index sh0, Index sh1) -> Code r a)
-> Val (Index sh0) -> Val (Index sh1) -> Code r a
forall a b c. (T (a, b) -> c) -> T a -> T b -> c
MultiValue.curry T (Index sh0, Index sh1) -> Code r a
Val (Index (sh0, sh1)) -> Code r a
forall r. Val (Index (sh0, sh1)) -> Code r a
code


fold1All ::
   (Shape.C sh, MultiValue.C a) =>
   (Exp a -> Exp a -> Exp a) ->
   Array sh a -> Exp a
fold1All :: forall sh a.
(C sh, C a) =>
(Exp a -> Exp a -> Exp a) -> Array sh a -> Exp a
fold1All Exp a -> Exp a -> Exp a
f (Array Exp sh
sh forall r. Val (Index sh) -> Code r a
code) = (forall r. CodeGenFunction r (T a)) -> Exp a
forall a. (forall r. CodeGenFunction r (T a)) -> Exp a
Exp ((Exp a -> Exp a -> Exp a)
-> Exp sh -> (Val (Index sh) -> Code r a) -> Code r a
forall sh ix a r.
(C sh, Index sh ~ ix, C a) =>
(Exp a -> Exp a -> Exp a)
-> Exp sh -> (Val ix -> Code r a) -> Code r a
fold1Code Exp a -> Exp a -> Exp a
f Exp sh
sh Val (Index sh) -> Code r a
forall r. Val (Index sh) -> Code r a
code)


findAllCode ::
   (Shape.C sh, Shape.Index sh ~ ix, MultiValue.C a) =>
   (Exp a -> Exp Bool) ->
   Exp sh ->
   (Val ix -> Code r a) ->
   Code r (Maybe a)
findAllCode :: forall sh ix a r.
(C sh, Index sh ~ ix, C a) =>
(Exp a -> Exp Bool)
-> Exp sh -> (Val ix -> Code r a) -> Code r (Maybe a)
findAllCode Exp a -> Exp Bool
p (Exp forall r. CodeGenFunction r (T sh)
sh) Val ix -> Code r a
code = do
   T sh
n <- CodeGenFunction r (T sh)
forall r. CodeGenFunction r (T sh)
sh
   T (Val a)
finalFound <-
      (Val a -> T (Val a) -> CodeGenFunction r (Value Bool, T (Val a)))
-> T r (Val a) -> T (Val a) -> CodeGenFunction r (T (Val a))
forall t a r.
Phi t =>
(a -> t -> CodeGenFunction r (Value Bool, t))
-> T r a -> t -> CodeGenFunction r t
Iter.mapWhileState_
         (\Val a
a T (Val a)
_found -> do
            MultiValue.Cons Repr Bool
b <- (Exp a -> Exp Bool) -> Val a -> CodeGenFunction r (T Bool)
forall ae am be bm r.
(Aggregate ae am, Aggregate be bm) =>
(ae -> be) -> am -> CodeGenFunction r bm
Expr.unliftM1 Exp a -> Exp Bool
p Val a
a
            Value Bool
notb <- Value Bool -> CodeGenFunction r (Value Bool)
forall (value :: * -> *) a r.
(ValueCons value, IsInteger a) =>
value a -> CodeGenFunction r (value a)
LLVM.inv Repr Bool
Value Bool
b
            (Value Bool, T (Val a))
-> CodeGenFunction r (Value Bool, T (Val a))
forall a. a -> CodeGenFunction r a
forall (m :: * -> *) a. Monad m => a -> m a
return (Value Bool
notb, Value Bool -> Val a -> T (Val a)
forall a. Value Bool -> a -> T a
Maybe.fromBool Repr Bool
Value Bool
b Val a
a))
         ((Val ix -> Code r a) -> T r (Val ix) -> T r (Val a)
forall a r b. (a -> CodeGenFunction r b) -> T r a -> T r b
Iter.mapM Val ix -> Code r a
code (T r (Val ix) -> T r (Val a)) -> T r (Val ix) -> T r (Val a)
forall a b. (a -> b) -> a -> b
$ T sh -> T r (Val ix)
forall sh ix r. (C sh, Index sh ~ ix) => T sh -> T r (T ix)
forall ix r. (Index sh ~ ix) => T sh -> T r (T ix)
Shape.iterator T sh
n)
         T (Val a)
forall a. Undefined a => T a
Maybe.nothing
   T (Val a)
-> Code r (Maybe a)
-> (Val a -> Code r (Maybe a))
-> Code r (Maybe a)
forall b a r.
Phi b =>
T a
-> CodeGenFunction r b
-> (a -> CodeGenFunction r b)
-> CodeGenFunction r b
Maybe.run T (Val a)
finalFound
      (Val (Maybe a) -> Code r (Maybe a)
forall a. a -> CodeGenFunction r a
forall (m :: * -> *) a. Monad m => a -> m a
return Val (Maybe a)
forall a. C a => T (Maybe a)
MultiValue.nothing)
      (Val (Maybe a) -> Code r (Maybe a)
forall a. a -> CodeGenFunction r a
forall (m :: * -> *) a. Monad m => a -> m a
return (Val (Maybe a) -> Code r (Maybe a))
-> (Val a -> Val (Maybe a)) -> Val a -> Code r (Maybe a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Val a -> Val (Maybe a)
forall a. T a -> T (Maybe a)
MultiValue.just)

{- |
In principle this can be implemented using fold1All
but it has a short-cut semantics.
@All@ means that it scans all dimensions
but it does not mean that it finds all occurrences.
If you want to get the index of the found element,
please decorate the array elements with their indices before calling 'findAll'.
-}
findAll ::
   (Shape.C sh, MultiValue.C a) =>
   (Exp a -> Exp Bool) ->
   Array sh a -> Exp (Maybe a)
findAll :: forall sh a.
(C sh, C a) =>
(Exp a -> Exp Bool) -> Array sh a -> Exp (Maybe a)
findAll Exp a -> Exp Bool
p (Array Exp sh
sh forall r. Val (Index sh) -> Code r a
code) = (forall r. CodeGenFunction r (T (Maybe a))) -> Exp (Maybe a)
forall a. (forall r. CodeGenFunction r (T a)) -> Exp a
Exp ((Exp a -> Exp Bool)
-> Exp sh -> (Val (Index sh) -> Code r a) -> Code r (Maybe a)
forall sh ix a r.
(C sh, Index sh ~ ix, C a) =>
(Exp a -> Exp Bool)
-> Exp sh -> (Val ix -> Code r a) -> Code r (Maybe a)
findAllCode Exp a -> Exp Bool
p Exp sh
sh Val (Index sh) -> Code r a
forall r. Val (Index sh) -> Code r a
code)


class Process proc where


infixl 3 $:.

{- |
Use this for combining several dimension manipulators.
E.g.

> apply (passAny $:. pick 3 $:. pass $:. replicate 10) array

The constraint @(Process proc0, Process proc1)@ is a bit weak.
We like to enforce that the type constructor like @Slice.T@
is the same in @proc0@ and @proc1@, and only the parameters differ.
Currently this coherence is achieved,
because we only provide functions of type @proc0 -> proc1@ with this condition.
-}
($:.) :: (Process proc0, Process proc1) => proc0 -> (proc0 -> proc1) -> proc1
$:. :: forall proc0 proc1.
(Process proc0, Process proc1) =>
proc0 -> (proc0 -> proc1) -> proc1
($:.) = ((proc0 -> proc1) -> proc0 -> proc1)
-> proc0 -> (proc0 -> proc1) -> proc1
forall a b c. (a -> b -> c) -> b -> a -> c
flip (proc0 -> proc1) -> proc0 -> proc1
forall a b. (a -> b) -> a -> b
($)