{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE EmptyDataDecls #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}
module Data.Array.Knead.Shape.Cubic (
   constant,
   paramWith,
   tunnel,

   T(..),
   Z(Z), z,
   (:.)((:.)),
   Shape,
   Index,
   cons, (#:.),
   head,
   tail,
   switchR,
   ) where

import qualified Data.Array.Knead.Shape as Shape
import qualified Data.Array.Knead.Shape.Cubic.Int as Index

import qualified Data.Array.Knead.Expression as Expr
import Data.Array.Knead.Expression (Exp, )

import qualified Data.Array.Comfort.Shape as ComfortShape
import Data.Array.Comfort.Shape (ZeroBased(ZeroBased))

import qualified LLVM.DSL.Parameter as Param

import qualified LLVM.Extra.Multi.Value.Marshal as Marshal
import qualified LLVM.Extra.Multi.Value as MultiValue
import qualified LLVM.Extra.Multi.Iterator as IterMV
import qualified LLVM.Extra.Iterator as Iter
import qualified LLVM.Extra.Arithmetic as A
import qualified LLVM.Extra.Tuple as Tuple
import qualified LLVM.Extra.Control as C
import LLVM.Extra.Multi.Value (Atom)

import qualified LLVM.Core as LLVM

import qualified Foreign.Storable as St
import Foreign.Storable.FixedArray (sizeOfArray, )
import Foreign.Ptr (castPtr, )

import qualified Type.Data.Num.Decimal as Dec
import qualified Type.Data.Num.Unary as Unary
import Type.Base.Proxy (Proxy(Proxy))

import qualified Data.Traversable as Trav
import qualified Data.Foldable as Fold
import qualified Data.FixedLength as FixedLength
import Data.FixedLength ((!:))

import Control.Monad (liftM2, )
import Control.Applicative (pure, (<$>), )

import Prelude hiding (min, head, tail, )


newtype T tag rank = Cons {forall tag rank. T tag rank -> T rank Int
decons :: FixedLength.T rank Index.Int}

data ShapeTag
data IndexTag

type Shape = T ShapeTag
type Index = T IndexTag


paramWith ::
   (Unary.Natural rank,
    Dec.Natural (Dec.FromUnary rank),
    Dec.Natural (Dec.FromUnary rank Dec.:*: LLVM.SizeOf Shape.Size)) =>
   Param.T p (T tag rank) ->
   (forall parameters.
    (Marshal.C parameters) =>
    (p -> parameters) ->
    (forall val. (Expr.Value val) =>
     MultiValue.T parameters -> val (T tag rank)) ->
    a) ->
   a
paramWith :: forall rank p tag a.
(Natural rank, Natural (FromUnary rank),
 Natural (FromUnary rank :*: SizeOf Size)) =>
T p (T tag rank)
-> (forall parameters.
    C parameters =>
    (p -> parameters)
    -> (forall (val :: * -> *).
        Value val =>
        T parameters -> val (T tag rank))
    -> a)
-> a
paramWith T p (T tag rank)
p forall parameters.
C parameters =>
(p -> parameters)
-> (forall (val :: * -> *).
    Value val =>
    T parameters -> val (T tag rank))
-> a
f =
   case T p (T tag rank) -> Tunnel p (T tag rank)
forall rank p tag.
(Natural rank, Natural (FromUnary rank),
 Natural (FromUnary rank :*: SizeOf Size)) =>
T p (T tag rank) -> Tunnel p (T tag rank)
tunnel T p (T tag rank)
p of
      Param.Tunnel p -> t
get T t -> T (T tag rank)
val -> (p -> t)
-> (forall (val :: * -> *). Value val => T t -> val (T tag rank))
-> a
forall parameters.
C parameters =>
(p -> parameters)
-> (forall (val :: * -> *).
    Value val =>
    T parameters -> val (T tag rank))
-> a
f p -> t
get (T (T tag rank) -> val (T tag rank)
forall a. T a -> val a
forall (val :: * -> *) a. Value val => T a -> val a
Expr.lift0 (T (T tag rank) -> val (T tag rank))
-> (T t -> T (T tag rank)) -> T t -> val (T tag rank)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. T t -> T (T tag rank)
val)

tunnel ::
   (Unary.Natural rank,
    Dec.Natural (Dec.FromUnary rank),
    Dec.Natural (Dec.FromUnary rank Dec.:*: LLVM.SizeOf Shape.Size)) =>
   Param.T p (T tag rank) -> Param.Tunnel p (T tag rank)
tunnel :: forall rank p tag.
(Natural rank, Natural (FromUnary rank),
 Natural (FromUnary rank :*: SizeOf Size)) =>
T p (T tag rank) -> Tunnel p (T tag rank)
tunnel T p (T tag rank)
p = (T tag rank -> T (T tag rank))
-> T p (T tag rank) -> Tunnel p (T tag rank)
forall a p. C a => (a -> T a) -> T p a -> Tunnel p a
Param.tunnel T tag rank -> T (T tag rank)
forall a. C a => a -> T a
MultiValue.cons T p (T tag rank)
p


data Z = Z
   deriving (Z -> Z -> Bool
(Z -> Z -> Bool) -> (Z -> Z -> Bool) -> Eq Z
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Z -> Z -> Bool
== :: Z -> Z -> Bool
$c/= :: Z -> Z -> Bool
/= :: Z -> Z -> Bool
Eq, Eq Z
Eq Z
-> (Z -> Z -> Ordering)
-> (Z -> Z -> Bool)
-> (Z -> Z -> Bool)
-> (Z -> Z -> Bool)
-> (Z -> Z -> Bool)
-> (Z -> Z -> Z)
-> (Z -> Z -> Z)
-> Ord Z
Z -> Z -> Bool
Z -> Z -> Ordering
Z -> Z -> Z
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: Z -> Z -> Ordering
compare :: Z -> Z -> Ordering
$c< :: Z -> Z -> Bool
< :: Z -> Z -> Bool
$c<= :: Z -> Z -> Bool
<= :: Z -> Z -> Bool
$c> :: Z -> Z -> Bool
> :: Z -> Z -> Bool
$c>= :: Z -> Z -> Bool
>= :: Z -> Z -> Bool
$cmax :: Z -> Z -> Z
max :: Z -> Z -> Z
$cmin :: Z -> Z -> Z
min :: Z -> Z -> Z
Ord, ReadPrec [Z]
ReadPrec Z
Int -> ReadS Z
ReadS [Z]
(Int -> ReadS Z)
-> ReadS [Z] -> ReadPrec Z -> ReadPrec [Z] -> Read Z
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
$creadsPrec :: Int -> ReadS Z
readsPrec :: Int -> ReadS Z
$creadList :: ReadS [Z]
readList :: ReadS [Z]
$creadPrec :: ReadPrec Z
readPrec :: ReadPrec Z
$creadListPrec :: ReadPrec [Z]
readListPrec :: ReadPrec [Z]
Read, Int -> Z -> ShowS
[Z] -> ShowS
Z -> String
(Int -> Z -> ShowS) -> (Z -> String) -> ([Z] -> ShowS) -> Show Z
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Z -> ShowS
showsPrec :: Int -> Z -> ShowS
$cshow :: Z -> String
show :: Z -> String
$cshowList :: [Z] -> ShowS
showList :: [Z] -> ShowS
Show)


infixl 3 :., #:.

data tail :. head = !tail :. !head
   deriving ((tail :. head) -> (tail :. head) -> Bool
((tail :. head) -> (tail :. head) -> Bool)
-> ((tail :. head) -> (tail :. head) -> Bool) -> Eq (tail :. head)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall tail head.
(Eq tail, Eq head) =>
(tail :. head) -> (tail :. head) -> Bool
$c== :: forall tail head.
(Eq tail, Eq head) =>
(tail :. head) -> (tail :. head) -> Bool
== :: (tail :. head) -> (tail :. head) -> Bool
$c/= :: forall tail head.
(Eq tail, Eq head) =>
(tail :. head) -> (tail :. head) -> Bool
/= :: (tail :. head) -> (tail :. head) -> Bool
Eq, Eq (tail :. head)
Eq (tail :. head)
-> ((tail :. head) -> (tail :. head) -> Ordering)
-> ((tail :. head) -> (tail :. head) -> Bool)
-> ((tail :. head) -> (tail :. head) -> Bool)
-> ((tail :. head) -> (tail :. head) -> Bool)
-> ((tail :. head) -> (tail :. head) -> Bool)
-> ((tail :. head) -> (tail :. head) -> tail :. head)
-> ((tail :. head) -> (tail :. head) -> tail :. head)
-> Ord (tail :. head)
(tail :. head) -> (tail :. head) -> Bool
(tail :. head) -> (tail :. head) -> Ordering
(tail :. head) -> (tail :. head) -> tail :. head
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall {tail} {head}. (Ord tail, Ord head) => Eq (tail :. head)
forall tail head.
(Ord tail, Ord head) =>
(tail :. head) -> (tail :. head) -> Bool
forall tail head.
(Ord tail, Ord head) =>
(tail :. head) -> (tail :. head) -> Ordering
forall tail head.
(Ord tail, Ord head) =>
(tail :. head) -> (tail :. head) -> tail :. head
$ccompare :: forall tail head.
(Ord tail, Ord head) =>
(tail :. head) -> (tail :. head) -> Ordering
compare :: (tail :. head) -> (tail :. head) -> Ordering
$c< :: forall tail head.
(Ord tail, Ord head) =>
(tail :. head) -> (tail :. head) -> Bool
< :: (tail :. head) -> (tail :. head) -> Bool
$c<= :: forall tail head.
(Ord tail, Ord head) =>
(tail :. head) -> (tail :. head) -> Bool
<= :: (tail :. head) -> (tail :. head) -> Bool
$c> :: forall tail head.
(Ord tail, Ord head) =>
(tail :. head) -> (tail :. head) -> Bool
> :: (tail :. head) -> (tail :. head) -> Bool
$c>= :: forall tail head.
(Ord tail, Ord head) =>
(tail :. head) -> (tail :. head) -> Bool
>= :: (tail :. head) -> (tail :. head) -> Bool
$cmax :: forall tail head.
(Ord tail, Ord head) =>
(tail :. head) -> (tail :. head) -> tail :. head
max :: (tail :. head) -> (tail :. head) -> tail :. head
$cmin :: forall tail head.
(Ord tail, Ord head) =>
(tail :. head) -> (tail :. head) -> tail :. head
min :: (tail :. head) -> (tail :. head) -> tail :. head
Ord, ReadPrec [tail :. head]
ReadPrec (tail :. head)
Int -> ReadS (tail :. head)
ReadS [tail :. head]
(Int -> ReadS (tail :. head))
-> ReadS [tail :. head]
-> ReadPrec (tail :. head)
-> ReadPrec [tail :. head]
-> Read (tail :. head)
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
forall tail head. (Read tail, Read head) => ReadPrec [tail :. head]
forall tail head. (Read tail, Read head) => ReadPrec (tail :. head)
forall tail head.
(Read tail, Read head) =>
Int -> ReadS (tail :. head)
forall tail head. (Read tail, Read head) => ReadS [tail :. head]
$creadsPrec :: forall tail head.
(Read tail, Read head) =>
Int -> ReadS (tail :. head)
readsPrec :: Int -> ReadS (tail :. head)
$creadList :: forall tail head. (Read tail, Read head) => ReadS [tail :. head]
readList :: ReadS [tail :. head]
$creadPrec :: forall tail head. (Read tail, Read head) => ReadPrec (tail :. head)
readPrec :: ReadPrec (tail :. head)
$creadListPrec :: forall tail head. (Read tail, Read head) => ReadPrec [tail :. head]
readListPrec :: ReadPrec [tail :. head]
Read, Int -> (tail :. head) -> ShowS
[tail :. head] -> ShowS
(tail :. head) -> String
(Int -> (tail :. head) -> ShowS)
-> ((tail :. head) -> String)
-> ([tail :. head] -> ShowS)
-> Show (tail :. head)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall tail head.
(Show tail, Show head) =>
Int -> (tail :. head) -> ShowS
forall tail head. (Show tail, Show head) => [tail :. head] -> ShowS
forall tail head.
(Show tail, Show head) =>
(tail :. head) -> String
$cshowsPrec :: forall tail head.
(Show tail, Show head) =>
Int -> (tail :. head) -> ShowS
showsPrec :: Int -> (tail :. head) -> ShowS
$cshow :: forall tail head.
(Show tail, Show head) =>
(tail :. head) -> String
show :: (tail :. head) -> String
$cshowList :: forall tail head. (Show tail, Show head) => [tail :. head] -> ShowS
showList :: [tail :. head] -> ShowS
Show)


(#:.) ::
   (Expr.Value val) =>
   val (T tag rank) -> val Index.Int -> val (T tag (Unary.Succ rank))
#:. :: forall (val :: * -> *) tag rank.
Value val =>
val (T tag rank) -> val Int -> val (T tag (Succ rank))
(#:.) = val (T tag rank) -> val Int -> val (T tag (Succ rank))
forall (val :: * -> *) tag rank.
Value val =>
val (T tag rank) -> val Int -> val (T tag (Succ rank))
cons

cons ::
   (Expr.Value val) =>
   val (T tag rank) -> val Index.Int -> val (T tag (Unary.Succ rank))
cons :: forall (val :: * -> *) tag rank.
Value val =>
val (T tag rank) -> val Int -> val (T tag (Succ rank))
cons =
   (T (T tag rank) -> T Int -> T (T tag (Succ rank)))
-> val (T tag rank) -> val Int -> val (T tag (Succ rank))
forall a b c. (T a -> T b -> T c) -> val a -> val b -> val c
forall (val :: * -> *) a b c.
Value val =>
(T a -> T b -> T c) -> val a -> val b -> val c
Expr.lift2 ((T (T tag rank) -> T Int -> T (T tag (Succ rank)))
 -> val (T tag rank) -> val Int -> val (T tag (Succ rank)))
-> (T (T tag rank) -> T Int -> T (T tag (Succ rank)))
-> val (T tag rank)
-> val Int
-> val (T tag (Succ rank))
forall a b. (a -> b) -> a -> b
$
      \(MultiValue.Cons Repr (T tag rank)
t) (MultiValue.Cons Repr Int
h) -> Repr (T tag (Succ rank)) -> T (T tag (Succ rank))
forall a. Repr a -> T a
MultiValue.Cons (Repr Int
Value Size
hValue Size -> T rank (Value Size) -> T (Succ rank) (Value Size)
forall a n. a -> T n a -> T (Succ n) a
!:T rank (Value Size)
Repr (T tag rank)
t)

z :: (Expr.Value val) => val (T tag Unary.Zero)
z :: forall (val :: * -> *) tag. Value val => val (T tag Zero)
z = T (T tag Zero) -> val (T tag Zero)
forall a. T a -> val a
forall (val :: * -> *) a. Value val => T a -> val a
Expr.lift0 (T (T tag Zero) -> val (T tag Zero))
-> T (T tag Zero) -> val (T tag Zero)
forall a b. (a -> b) -> a -> b
$ Repr (T tag Zero) -> T (T tag Zero)
forall a. Repr a -> T a
MultiValue.Cons T Zero (Value Size)
Repr (T tag Zero)
forall a. T Zero a
FixedLength.end

head ::
   (Expr.Value val, Unary.Natural rank) =>
   val (T tag (Unary.Succ rank)) -> val Index.Int
head :: forall (val :: * -> *) rank tag.
(Value val, Natural rank) =>
val (T tag (Succ rank)) -> val Int
head =
   (T (T tag (Succ rank)) -> T Int)
-> val (T tag (Succ rank)) -> val Int
forall a b. (T a -> T b) -> val a -> val b
forall (val :: * -> *) a b.
Value val =>
(T a -> T b) -> val a -> val b
Expr.lift1 ((T (T tag (Succ rank)) -> T Int)
 -> val (T tag (Succ rank)) -> val Int)
-> (T (T tag (Succ rank)) -> T Int)
-> val (T tag (Succ rank))
-> val Int
forall a b. (a -> b) -> a -> b
$ \(MultiValue.Cons Repr (T tag (Succ rank))
sh) -> Repr Int -> T Int
forall a. Repr a -> T a
MultiValue.Cons (Repr Int -> T Int) -> Repr Int -> T Int
forall a b. (a -> b) -> a -> b
$ T (Succ rank) (Value Size) -> Value Size
forall n a. Positive n => T n a -> a
FixedLength.head T (Succ rank) (Value Size)
Repr (T tag (Succ rank))
sh

tail ::
   (Expr.Value val, Unary.Natural rank) =>
   val (T tag (Unary.Succ rank)) -> val (T tag rank)
tail :: forall (val :: * -> *) rank tag.
(Value val, Natural rank) =>
val (T tag (Succ rank)) -> val (T tag rank)
tail =
   (T (T tag (Succ rank)) -> T (T tag rank))
-> val (T tag (Succ rank)) -> val (T tag rank)
forall a b. (T a -> T b) -> val a -> val b
forall (val :: * -> *) a b.
Value val =>
(T a -> T b) -> val a -> val b
Expr.lift1 ((T (T tag (Succ rank)) -> T (T tag rank))
 -> val (T tag (Succ rank)) -> val (T tag rank))
-> (T (T tag (Succ rank)) -> T (T tag rank))
-> val (T tag (Succ rank))
-> val (T tag rank)
forall a b. (a -> b) -> a -> b
$ \(MultiValue.Cons Repr (T tag (Succ rank))
sh) -> Repr (T tag rank) -> T (T tag rank)
forall a. Repr a -> T a
MultiValue.Cons (Repr (T tag rank) -> T (T tag rank))
-> Repr (T tag rank) -> T (T tag rank)
forall a b. (a -> b) -> a -> b
$ T (Succ rank) (Value Size) -> T rank (Value Size)
forall n a. T (Succ n) a -> T n a
FixedLength.tail T (Succ rank) (Value Size)
Repr (T tag (Succ rank))
sh

switchR ::
   (Unary.Natural rank) =>
   Expr.Value val =>
   (val (T tag rank) -> val Index.Int -> a) ->
   val (T tag (Unary.Succ rank)) -> a
switchR :: forall rank (val :: * -> *) tag a.
(Natural rank, Value val) =>
(val (T tag rank) -> val Int -> a) -> val (T tag (Succ rank)) -> a
switchR val (T tag rank) -> val Int -> a
f val (T tag (Succ rank))
ix = val (T tag rank) -> val Int -> a
f (val (T tag (Succ rank)) -> val (T tag rank)
forall (val :: * -> *) rank tag.
(Value val, Natural rank) =>
val (T tag (Succ rank)) -> val (T tag rank)
tail val (T tag (Succ rank))
ix) (val (T tag (Succ rank)) -> val Int
forall (val :: * -> *) rank tag.
(Value val, Natural rank) =>
val (T tag (Succ rank)) -> val Int
head val (T tag (Succ rank))
ix)


rank :: T tag rank -> Proxy rank
rank :: forall tag rank. T tag rank -> Proxy rank
rank (Cons T rank Int
_) = Proxy rank
forall a. Proxy a
Proxy


instance (tag ~ ShapeTag, rank ~ Unary.Zero) => Shape.Scalar (T tag rank) where
   scalar :: forall (val :: * -> *). Value val => val (T tag rank)
scalar = T (T tag rank) -> val (T tag rank)
forall a. T a -> val a
forall (val :: * -> *) a. Value val => T a -> val a
Expr.lift0 (T (T tag rank) -> val (T tag rank))
-> T (T tag rank) -> val (T tag rank)
forall a b. (a -> b) -> a -> b
$ Repr (T tag rank) -> T (T tag rank)
forall a. Repr a -> T a
MultiValue.Cons T Zero (Value Size)
Repr (T tag rank)
forall a. T Zero a
FixedLength.end
   zeroIndex :: forall (val :: * -> *) (f :: * -> *).
Value val =>
f (T tag rank) -> val (Index (T tag rank))
zeroIndex f (T tag rank)
_ = T (Index (T tag rank)) -> val (Index (T tag rank))
forall a. T a -> val a
forall (val :: * -> *) a. Value val => T a -> val a
Expr.lift0 (T (Index (T tag rank)) -> val (Index (T tag rank)))
-> T (Index (T tag rank)) -> val (Index (T tag rank))
forall a b. (a -> b) -> a -> b
$ Repr (Index Zero) -> T (Index Zero)
forall a. Repr a -> T a
MultiValue.Cons T Zero (Value Size)
Repr (Index Zero)
forall a. T Zero a
FixedLength.end


type family AtomRank sh
type instance AtomRank (Atom (T tag rank)) = rank
type instance AtomRank (sh:.s) = Unary.Succ (AtomRank s)

type family AtomTag sh
type instance AtomTag (Atom (T tag rank)) = tag
type instance AtomTag (sh:.s) = AtomTag sh

type instance MultiValue.PatternTuple (sh:.s) =
   T (AtomTag sh) (Unary.Succ (AtomRank sh))

type instance MultiValue.Decomposed f (sh:.s) =
   MultiValue.Decomposed f sh :. f Index.Int

instance
   (Expr.Decompose sh, Expr.Decompose s,
    MultiValue.Decomposed Exp s ~ Exp Index.Int,
    MultiValue.PatternTuple s ~ Index.Int,
    MultiValue.PatternTuple sh ~ T (AtomTag sh) (AtomRank sh),
    Unary.Natural (AtomRank sh)) =>
      Expr.Decompose (sh :. s) where
   decompose :: (sh :. s)
-> Exp (PatternTuple (sh :. s)) -> Decomposed Exp (sh :. s)
decompose (sh
psh:.s
ps) Exp (PatternTuple (sh :. s))
x =
      sh -> Exp (PatternTuple sh) -> Decomposed Exp sh
forall pattern.
Decompose pattern =>
pattern -> Exp (PatternTuple pattern) -> Decomposed Exp pattern
Expr.decompose sh
psh (Exp (T (AtomTag sh) (Succ (AtomRank sh)))
-> Exp (T (AtomTag sh) (AtomRank sh))
forall (val :: * -> *) rank tag.
(Value val, Natural rank) =>
val (T tag (Succ rank)) -> val (T tag rank)
tail Exp (PatternTuple (sh :. s))
Exp (T (AtomTag sh) (Succ (AtomRank sh)))
x) Decomposed Exp sh -> Exp Int -> Decomposed Exp sh :. Exp Int
forall tail head. tail -> head -> tail :. head
:. s -> Exp (PatternTuple s) -> Decomposed Exp s
forall pattern.
Decompose pattern =>
pattern -> Exp (PatternTuple pattern) -> Decomposed Exp pattern
Expr.decompose s
ps (Exp (T (AtomTag sh) (Succ (AtomRank sh))) -> Exp Int
forall (val :: * -> *) rank tag.
(Value val, Natural rank) =>
val (T tag (Succ rank)) -> val Int
head Exp (PatternTuple (sh :. s))
Exp (T (AtomTag sh) (Succ (AtomRank sh)))
x)


type family Rank sh
type instance Rank (T tag rank) = rank

type family Tag sh
type instance Tag (T tag rank) = tag

instance
   (Expr.Compose sh,
    Expr.Composed sh ~ T (Tag (Expr.Composed sh)) (Rank (Expr.Composed sh)),
    Expr.Compose s,
    Expr.Composed s ~ Index.Int) =>
      Expr.Compose (sh :. s) where
   type Composed (sh :. s) =
            T (Tag (Expr.Composed sh)) (Unary.Succ (Rank (Expr.Composed sh)))
   compose :: (sh :. s) -> Exp (Composed (sh :. s))
compose (sh
sh :. s
s) = Exp (T (Tag (Composed sh)) (Rank (Composed sh)))
-> Exp Int
-> Exp (T (Tag (Composed sh)) (Succ (Rank (Composed sh))))
forall (val :: * -> *) tag rank.
Value val =>
val (T tag rank) -> val Int -> val (T tag (Succ rank))
cons (sh -> Exp (Composed sh)
forall multituple.
Compose multituple =>
multituple -> Exp (Composed multituple)
Expr.compose sh
sh) (s -> Exp (Composed s)
forall multituple.
Compose multituple =>
multituple -> Exp (Composed multituple)
Expr.compose s
s)


instance (Unary.Natural rank) => St.Storable (T tag rank) where
   sizeOf :: T tag rank -> Int
sizeOf T tag rank
sh = Int -> Size -> Int
forall a. Storable a => Int -> a -> Int
sizeOfArray (Proxy rank -> Int
forall n a. (Natural n, Num a) => Proxy n -> a
Unary.integralFromProxy (Proxy rank -> Int) -> Proxy rank -> Int
forall a b. (a -> b) -> a -> b
$ T tag rank -> Proxy rank
forall tag rank. T tag rank -> Proxy rank
rank T tag rank
sh) (Size
0::Shape.Size)
   alignment :: T tag rank -> Int
alignment (Cons T rank Int
_sh) = Size -> Int
forall a. Storable a => a -> Int
St.alignment (Size
0::Shape.Size)
   poke :: Ptr (T tag rank) -> T tag rank -> IO ()
poke Ptr (T tag rank)
ptr = Ptr (T rank Size) -> T rank Size -> IO ()
forall a. Storable a => Ptr a -> a -> IO ()
St.poke (Ptr (T tag rank) -> Ptr (T rank Size)
forall a b. Ptr a -> Ptr b
castPtr Ptr (T tag rank)
ptr) (T rank Size -> IO ())
-> (T tag rank -> T rank Size) -> T tag rank -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Size) -> T rank Int -> T rank Size
forall a b. (a -> b) -> T rank a -> T rank b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(Index.Int Size
i) -> Size
i) (T rank Int -> T rank Size)
-> (T tag rank -> T rank Int) -> T tag rank -> T rank Size
forall b c a. (b -> c) -> (a -> b) -> a -> c
. T tag rank -> T rank Int
forall tag rank. T tag rank -> T rank Int
decons
   peek :: Ptr (T tag rank) -> IO (T tag rank)
peek = (T rank Size -> T tag rank) -> IO (T rank Size) -> IO (T tag rank)
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (T rank Int -> T tag rank
forall tag rank. T rank Int -> T tag rank
Cons (T rank Int -> T tag rank)
-> (T rank Size -> T rank Int) -> T rank Size -> T tag rank
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Size -> Int) -> T rank Size -> T rank Int
forall a b. (a -> b) -> T rank a -> T rank b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Size -> Int
Index.Int) (IO (T rank Size) -> IO (T tag rank))
-> (Ptr (T tag rank) -> IO (T rank Size))
-> Ptr (T tag rank)
-> IO (T tag rank)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ptr (T rank Size) -> IO (T rank Size)
forall a. Storable a => Ptr a -> IO a
St.peek (Ptr (T rank Size) -> IO (T rank Size))
-> (Ptr (T tag rank) -> Ptr (T rank Size))
-> Ptr (T tag rank)
-> IO (T rank Size)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ptr (T tag rank) -> Ptr (T rank Size)
forall a b. Ptr a -> Ptr b
castPtr

instance
   (Unary.Natural rank,
    Dec.Natural (Dec.FromUnary rank),
    Dec.Natural (Dec.FromUnary rank Dec.:*: LLVM.SizeOf Shape.Size)) =>
      Marshal.C (T tag rank) where
   pack :: T tag rank -> Struct (T tag rank)
pack = [Size] -> Array (FromUnary rank) Size
forall n a. [a] -> Array n a
LLVM.Array ([Size] -> Array (FromUnary rank) Size)
-> (T tag rank -> [Size])
-> T tag rank
-> Array (FromUnary rank) Size
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Size) -> [Int] -> [Size]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Size
Int -> Struct Int
forall a. C a => a -> Struct a
Marshal.pack ([Int] -> [Size]) -> (T tag rank -> [Int]) -> T tag rank -> [Size]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. T rank Int -> [Int]
forall a. T rank a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
Fold.toList (T rank Int -> [Int])
-> (T tag rank -> T rank Int) -> T tag rank -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. T tag rank -> T rank Int
forall tag rank. T tag rank -> T rank Int
decons
   unpack :: Struct (T tag rank) -> T tag rank
unpack (LLVM.Array [Size]
sh) = T rank Int -> T tag rank
forall tag rank. T rank Int -> T tag rank
Cons (T rank Int -> T tag rank) -> T rank Int -> T tag rank
forall a b. (a -> b) -> a -> b
$ [Int] -> T rank Int
forall n a. Natural n => [a] -> T n a
toFixedList ([Int] -> T rank Int) -> [Int] -> T rank Int
forall a b. (a -> b) -> a -> b
$ (Size -> Int) -> [Size] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Size -> Int
Struct Int -> Int
forall a. C a => Struct a -> a
Marshal.unpack [Size]
sh

toFixedList :: (Unary.Natural n) => [a] -> FixedLength.T n a
toFixedList :: forall n a. Natural n => [a] -> T n a
toFixedList [a]
xs = ([a], T n a) -> T n a
forall a b. (a, b) -> b
snd (([a], T n a) -> T n a) -> ([a], T n a) -> T n a
forall a b. (a -> b) -> a -> b
$ ([a] -> () -> ([a], a)) -> [a] -> T n () -> ([a], T n a)
forall (t :: * -> *) s a b.
Traversable t =>
(s -> a -> (s, b)) -> s -> t a -> (s, t b)
Trav.mapAccumL (\(a
y:[a]
ys) () -> ([a]
ys,a
y)) [a]
xs (() -> T n ()
forall a. a -> T n a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ())


instance (Unary.Natural rank) => MultiValue.C (T tag rank) where
   type Repr (T tag rank) = FixedLength.T rank (MultiValue.Repr Index.Int)
   cons :: T tag rank -> T (T tag rank)
cons = T rank (Value Size) -> T (T tag rank)
Repr (T tag rank) -> T (T tag rank)
forall a. Repr a -> T a
MultiValue.Cons (T rank (Value Size) -> T (T tag rank))
-> (T tag rank -> T rank (Value Size))
-> T tag rank
-> T (T tag rank)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Value Size) -> T rank Int -> T rank (Value Size)
forall a b. (a -> b) -> T rank a -> T rank b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(Index.Int Size
i) -> Size -> Value Size
forall a. IsConst a => a -> Value a
LLVM.valueOf Size
i) (T rank Int -> T rank (Value Size))
-> (T tag rank -> T rank Int) -> T tag rank -> T rank (Value Size)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. T tag rank -> T rank Int
forall tag rank. T tag rank -> T rank Int
decons
   undef :: T (T tag rank)
undef = T Int -> T (T tag rank)
forall rank tag. Natural rank => T Int -> T (T tag rank)
constant (T Int -> T (T tag rank)) -> T Int -> T (T tag rank)
forall a b. (a -> b) -> a -> b
$ T Int
forall a. C a => T a
MultiValue.undef
   zero :: T (T tag rank)
zero = T Int -> T (T tag rank)
forall rank tag. Natural rank => T Int -> T (T tag rank)
constant (T Int -> T (T tag rank)) -> T Int -> T (T tag rank)
forall a b. (a -> b) -> a -> b
$ T Int
forall a. C a => T a
MultiValue.zero
   addPhi :: forall r.
BasicBlock
-> T (T tag rank) -> T (T tag rank) -> CodeGenFunction r ()
addPhi BasicBlock
bb (MultiValue.Cons Repr (T tag rank)
a) (MultiValue.Cons Repr (T tag rank)
b) =
      BasicBlock
-> T rank (Value Size)
-> T rank (Value Size)
-> CodeGenFunction r ()
forall a (f :: * -> *) r.
(Phi a, Foldable f, Applicative f) =>
BasicBlock -> f a -> f a -> CodeGenFunction r ()
Tuple.addPhiFoldable BasicBlock
bb T rank (Value Size)
Repr (T tag rank)
a T rank (Value Size)
Repr (T tag rank)
b
   phi :: forall r.
BasicBlock -> T (T tag rank) -> CodeGenFunction r (T (T tag rank))
phi BasicBlock
bb (MultiValue.Cons Repr (T tag rank)
a) =
      (T rank (Value Size) -> T (T tag rank))
-> CodeGenFunction r (T rank (Value Size))
-> CodeGenFunction r (T (T tag rank))
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 rank (Value Size) -> T (T tag rank)
Repr (T tag rank) -> T (T tag rank)
forall a. Repr a -> T a
MultiValue.Cons (CodeGenFunction r (T rank (Value Size))
 -> CodeGenFunction r (T (T tag rank)))
-> (Repr (T tag rank) -> CodeGenFunction r (T rank (Value Size)))
-> Repr (T tag rank)
-> CodeGenFunction r (T (T tag rank))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BasicBlock
-> T rank (Value Size) -> CodeGenFunction r (T rank (Value Size))
forall a (f :: * -> *) r.
(Phi a, Traversable f) =>
BasicBlock -> f a -> CodeGenFunction r (f a)
Tuple.phiTraversable BasicBlock
bb (Repr (T tag rank) -> CodeGenFunction r (T (T tag rank)))
-> Repr (T tag rank) -> CodeGenFunction r (T (T tag rank))
forall a b. (a -> b) -> a -> b
$ Repr (T tag rank)
a

constant ::
   (Unary.Natural rank) => MultiValue.T Index.Int -> MultiValue.T (T tag rank)
constant :: forall rank tag. Natural rank => T Int -> T (T tag rank)
constant (MultiValue.Cons Repr Int
x) = Repr (T tag rank) -> T (T tag rank)
forall a. Repr a -> T a
MultiValue.Cons (Repr (T tag rank) -> T (T tag rank))
-> Repr (T tag rank) -> T (T tag rank)
forall a b. (a -> b) -> a -> b
$ Value Size -> T rank (Value Size)
forall a. a -> T rank a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Repr Int
Value Size
x

instance
   (tag ~ ShapeTag, Unary.Natural rank) =>
      ComfortShape.C (T tag rank) where
   size :: T tag rank -> Int
size = T rank Int -> Int
forall a. Num a => T rank a -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
Fold.product (T rank Int -> Int)
-> (T tag rank -> T rank Int) -> T tag rank -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int) -> T rank Int -> T rank Int
forall a b. (a -> b) -> T rank a -> T rank b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (ZeroBased Size -> Int
forall sh. C sh => sh -> Int
ComfortShape.size (ZeroBased Size -> Int) -> (Int -> ZeroBased Size) -> Int -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> ZeroBased Size
shapeFromInt) (T rank Int -> T rank Int)
-> (T tag rank -> T rank Int) -> T tag rank -> T rank Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. T tag rank -> T rank Int
forall tag rank. T tag rank -> T rank Int
decons

instance
   (tag ~ ShapeTag, Unary.Natural rank) =>
      ComfortShape.Indexed (T tag rank) where
   type Index (T tag rank) = Index rank
   indices :: T tag rank -> [Index (T tag rank)]
indices (Cons T rank Int
ix) =
      (T rank Size -> Index (T tag rank))
-> [T rank Size] -> [Index (T tag rank)]
forall a b. (a -> b) -> [a] -> [b]
map (T rank Int -> T IndexTag rank
forall tag rank. T rank Int -> T tag rank
Cons (T rank Int -> T IndexTag rank)
-> (T rank Size -> T rank Int) -> T rank Size -> T IndexTag rank
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Size -> Int) -> T rank Size -> T rank Int
forall a b. (a -> b) -> T rank a -> T rank b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Size -> Int
Index.Int) ([T rank Size] -> [Index (T tag rank)])
-> [T rank Size] -> [Index (T tag rank)]
forall a b. (a -> b) -> a -> b
$
      (Int -> [Size]) -> T rank Int -> [T rank Size]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> T rank a -> m (T rank b)
Trav.mapM (ZeroBased Size -> [Size]
ZeroBased Size -> [Index (ZeroBased Size)]
forall sh. Indexed sh => sh -> [Index sh]
ComfortShape.indices (ZeroBased Size -> [Size])
-> (Int -> ZeroBased Size) -> Int -> [Size]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> ZeroBased Size
shapeFromInt) T rank Int
ix
   inBounds :: T tag rank -> Index (T tag rank) -> Bool
inBounds (Cons T rank Int
sh) (Cons T rank Int
ix) =
      T rank Bool -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
Fold.and (T rank Bool -> Bool) -> T rank Bool -> Bool
forall a b. (a -> b) -> a -> b
$
      (ZeroBased Size -> Size -> Bool)
-> T rank (ZeroBased Size) -> T rank Size -> T rank Bool
forall n a b c.
Natural n =>
(a -> b -> c) -> T n a -> T n b -> T n c
FixedLength.zipWith ZeroBased Size -> Size -> Bool
ZeroBased Size -> Index (ZeroBased Size) -> Bool
forall sh. Indexed sh => sh -> Index sh -> Bool
ComfortShape.inBounds
         (Int -> ZeroBased Size
shapeFromInt (Int -> ZeroBased Size) -> T rank Int -> T rank (ZeroBased Size)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> T rank Int
sh) (Int -> Size
indexFromInt (Int -> Size) -> T rank Int -> T rank Size
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> T rank Int
ix)
   unifiedOffset :: forall check.
Checking check =>
T tag rank -> Index (T tag rank) -> Result check Int
unifiedOffset (Cons T rank Int
sh) (Cons T rank Int
ix) =
      (Int -> (ZeroBased Size, Size) -> Result check Int)
-> Int -> T rank (ZeroBased Size, Size) -> Result check Int
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
Fold.foldlM
         (\Int
off (ZeroBased Size
s,Size
i) -> do
            Int
ioff <- ZeroBased Size -> Index (ZeroBased Size) -> Result check Int
forall sh check.
(Indexed sh, Checking check) =>
sh -> Index sh -> Result check Int
forall check.
Checking check =>
ZeroBased Size -> Index (ZeroBased Size) -> Result check Int
ComfortShape.unifiedOffset ZeroBased Size
s Size
Index (ZeroBased Size)
i
            Int -> Result check Int
forall a. a -> Result check a
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> Result check Int) -> Int -> Result check Int
forall a b. (a -> b) -> a -> b
$! Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
* ZeroBased Size -> Int
forall sh. C sh => sh -> Int
ComfortShape.size ZeroBased Size
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
ioff)
         Int
0 (T rank (ZeroBased Size, Size) -> Result check Int)
-> T rank (ZeroBased Size, Size) -> Result check Int
forall a b. (a -> b) -> a -> b
$
      (ZeroBased Size -> Size -> (ZeroBased Size, Size))
-> T rank (ZeroBased Size)
-> T rank Size
-> T rank (ZeroBased Size, Size)
forall n a b c.
Natural n =>
(a -> b -> c) -> T n a -> T n b -> T n c
FixedLength.zipWith (,) (Int -> ZeroBased Size
shapeFromInt (Int -> ZeroBased Size) -> T rank Int -> T rank (ZeroBased Size)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> T rank Int
sh) (Int -> Size
indexFromInt (Int -> Size) -> T rank Int -> T rank Size
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> T rank Int
ix)

shapeFromInt :: Index.Int -> ZeroBased Shape.Size
shapeFromInt :: Int -> ZeroBased Size
shapeFromInt (Index.Int Size
i) = Size -> ZeroBased Size
forall n. n -> ZeroBased n
ZeroBased Size
i

indexFromInt :: Index.Int -> Shape.Size
indexFromInt :: Int -> Size
indexFromInt (Index.Int Size
i) = Size
i


instance (tag ~ ShapeTag, Unary.Natural rank) => Shape.C (T tag rank) where
   size :: forall r. T (T tag rank) -> CodeGenFunction r (Value Size)
size (MultiValue.Cons Repr (T tag rank)
sh) = (Value Size -> Value Size -> CodeGenFunction r (Value Size))
-> Value Size
-> T rank (Value Size)
-> CodeGenFunction r (Value Size)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
Fold.foldlM Value Size -> Value Size -> CodeGenFunction r (Value Size)
forall a r. PseudoRing a => a -> a -> CodeGenFunction r a
forall r.
Value Size -> Value Size -> CodeGenFunction r (Value Size)
A.mul Value Size
forall a. IntegerConstant a => a
A.one T rank (Value Size)
Repr (T tag rank)
sh
   intersectCode :: forall r.
T (T tag rank)
-> T (T tag rank) -> CodeGenFunction r (T (T tag rank))
intersectCode (MultiValue.Cons Repr (T tag rank)
sh0) (MultiValue.Cons Repr (T tag rank)
sh1) =
      (T rank (Value Size) -> T (T tag rank))
-> CodeGenFunction r (T rank (Value Size))
-> CodeGenFunction r (T (T tag rank))
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 rank (Value Size) -> T (T tag rank)
Repr (T tag rank) -> T (T tag rank)
forall a. Repr a -> T a
MultiValue.Cons (CodeGenFunction r (T rank (Value Size))
 -> CodeGenFunction r (T (T tag rank)))
-> CodeGenFunction r (T rank (Value Size))
-> CodeGenFunction r (T (T tag rank))
forall a b. (a -> b) -> a -> b
$ T rank (CodeGenFunction r (Value Size))
-> CodeGenFunction r (T rank (Value Size))
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
forall (m :: * -> *) a. Monad m => T rank (m a) -> m (T rank a)
Trav.sequence (T rank (CodeGenFunction r (Value Size))
 -> CodeGenFunction r (T rank (Value Size)))
-> T rank (CodeGenFunction r (Value Size))
-> CodeGenFunction r (T rank (Value Size))
forall a b. (a -> b) -> a -> b
$ (Value Size -> Value Size -> CodeGenFunction r (Value Size))
-> T rank (Value Size)
-> T rank (Value Size)
-> T rank (CodeGenFunction r (Value Size))
forall n a b c.
Natural n =>
(a -> b -> c) -> T n a -> T n b -> T n c
FixedLength.zipWith Value Size -> Value Size -> CodeGenFunction r (Value Size)
forall a r. Real a => a -> a -> CodeGenFunction r a
forall r.
Value Size -> Value Size -> CodeGenFunction r (Value Size)
A.min T rank (Value Size)
Repr (T tag rank)
sh0 T rank (Value Size)
Repr (T tag rank)
sh1
   sizeOffset :: forall ix r.
(Index (T tag rank) ~ ix) =>
T (T tag rank)
-> CodeGenFunction
     r (Value Size, T ix -> CodeGenFunction r (Value Size))
sizeOffset T (T tag rank)
sh =
      -- would a joint implementation be more efficient?
      (Value Size
 -> (T ix -> CodeGenFunction r (Value Size))
 -> (Value Size, T ix -> CodeGenFunction r (Value Size)))
-> CodeGenFunction r (Value Size)
-> CodeGenFunction r (T ix -> CodeGenFunction r (Value Size))
-> CodeGenFunction
     r (Value Size, T ix -> CodeGenFunction r (Value Size))
forall (m :: * -> *) a1 a2 r.
Monad m =>
(a1 -> a2 -> r) -> m a1 -> m a2 -> m r
liftM2 (,) (T (T tag rank) -> CodeGenFunction r (Value Size)
forall r. T (T tag rank) -> CodeGenFunction r (Value Size)
forall sh r. C sh => T sh -> CodeGenFunction r (Value Size)
Shape.size T (T tag rank)
sh) ((T ix -> CodeGenFunction r (Value Size))
-> CodeGenFunction r (T ix -> CodeGenFunction r (Value Size))
forall a. a -> CodeGenFunction r a
forall (m :: * -> *) a. Monad m => a -> m a
return ((T ix -> CodeGenFunction r (Value Size))
 -> CodeGenFunction r (T ix -> CodeGenFunction r (Value Size)))
-> (T ix -> CodeGenFunction r (Value Size))
-> CodeGenFunction r (T ix -> CodeGenFunction r (Value Size))
forall a b. (a -> b) -> a -> b
$ T (T ShapeTag rank)
-> T (Index rank) -> CodeGenFunction r (Value Size)
forall rank r.
Natural rank =>
T (Shape rank) -> T (Index rank) -> CodeGenFunction r (Value Size)
offsetCode T (T tag rank)
T (T ShapeTag rank)
sh)
   iterator :: forall ix r.
(Index (T tag rank) ~ ix) =>
T (T tag rank) -> T r (T ix)
iterator = T (T tag rank) -> T r (T ix)
T (T ShapeTag rank) -> T r (T (Index rank))
forall rank r.
Natural rank =>
T (Shape rank) -> T r (T (Index rank))
iterator
   loop :: forall ix state r.
(Index (T tag rank) ~ ix, Phi state) =>
(T ix -> state -> CodeGenFunction r state)
-> T (T tag rank) -> state -> CodeGenFunction r state
loop = (T ix -> state -> CodeGenFunction r state)
-> T (T tag rank) -> state -> CodeGenFunction r state
(T (Index rank) -> state -> CodeGenFunction r state)
-> T (T ShapeTag rank) -> state -> CodeGenFunction r state
forall rank state r.
(Natural rank, Phi state) =>
(T (Index rank) -> state -> CodeGenFunction r state)
-> T (Shape rank) -> state -> CodeGenFunction r state
loop


offsetCode ::
   (Unary.Natural rank) =>
   MultiValue.T (Shape rank) -> MultiValue.T (Index rank) ->
   LLVM.CodeGenFunction r (LLVM.Value Shape.Size)
offsetCode :: forall rank r.
Natural rank =>
T (Shape rank) -> T (Index rank) -> CodeGenFunction r (Value Size)
offsetCode (MultiValue.Cons Repr (Shape rank)
sh) (MultiValue.Cons Repr (Index rank)
ix) =
   (Value Size
 -> (Value Size, Value Size) -> CodeGenFunction r (Value Size))
-> Value Size
-> T rank (Value Size, Value Size)
-> CodeGenFunction r (Value Size)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
Fold.foldlM (\Value Size
off (Value Size
s,Value Size
i) -> Value Size -> Value Size -> CodeGenFunction r (Value Size)
forall a r. PseudoRing a => a -> a -> CodeGenFunction r a
forall r.
Value Size -> Value Size -> CodeGenFunction r (Value Size)
A.mul Value Size
off Value Size
s CodeGenFunction r (Value Size)
-> (Value Size -> CodeGenFunction r (Value Size))
-> CodeGenFunction r (Value Size)
forall a b.
CodeGenFunction r a
-> (a -> CodeGenFunction r b) -> CodeGenFunction r b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Value Size -> Value Size -> CodeGenFunction r (Value Size)
forall a r. Additive a => a -> a -> CodeGenFunction r a
forall r.
Value Size -> Value Size -> CodeGenFunction r (Value Size)
A.add Value Size
i) Value Size
forall a. Additive a => a
A.zero (T rank (Value Size, Value Size) -> CodeGenFunction r (Value Size))
-> T rank (Value Size, Value Size)
-> CodeGenFunction r (Value Size)
forall a b. (a -> b) -> a -> b
$
   (Value Size -> Value Size -> (Value Size, Value Size))
-> T rank (Value Size)
-> T rank (Value Size)
-> T rank (Value Size, Value Size)
forall n a b c.
Natural n =>
(a -> b -> c) -> T n a -> T n b -> T n c
FixedLength.zipWith (,) T rank (Value Size)
Repr (Shape rank)
sh T rank (Value Size)
Repr (Index rank)
ix


newtype Iterator r rank =
   Iterator {
      forall r rank.
Iterator r rank -> T (Shape rank) -> T r (T (Index rank))
runIterator ::
         MultiValue.T (Shape rank) -> Iter.T r (MultiValue.T (Index rank))
   }

iterator ::
   (Unary.Natural rank) =>
   MultiValue.T (Shape rank) -> Iter.T r (MultiValue.T (Index rank))
iterator :: forall rank r.
Natural rank =>
T (Shape rank) -> T r (T (Index rank))
iterator =
   Iterator r rank -> T (Shape rank) -> T r (T (Index rank))
forall r rank.
Iterator r rank -> T (Shape rank) -> T r (T (Index rank))
runIterator (Iterator r rank -> T (Shape rank) -> T r (T (Index rank)))
-> Iterator r rank -> T (Shape rank) -> T r (T (Index rank))
forall a b. (a -> b) -> a -> b
$
   Iterator r Zero
-> (forall m. Natural m => Iterator r (Succ m)) -> Iterator r rank
forall n (f :: * -> *).
Natural n =>
f Zero -> (forall m. Natural m => f (Succ m)) -> f n
forall (f :: * -> *).
f Zero -> (forall m. Natural m => f (Succ m)) -> f rank
Unary.switchNat
      ((T (T ShapeTag Zero) -> T r (T (Index Zero))) -> Iterator r Zero
forall r rank.
(T (Shape rank) -> T r (T (Index rank))) -> Iterator r rank
Iterator ((T (T ShapeTag Zero) -> T r (T (Index Zero))) -> Iterator r Zero)
-> (T (T ShapeTag Zero) -> T r (T (Index Zero))) -> Iterator r Zero
forall a b. (a -> b) -> a -> b
$ \ T (T ShapeTag Zero)
_z -> T (Index Zero) -> T r (T (Index Zero))
forall a r. a -> T r a
Iter.singleton T (Index Zero)
forall (val :: * -> *) tag. Value val => val (T tag Zero)
z)
      ((T (Shape (Succ m)) -> T r (T (Index (Succ m))))
-> Iterator r (Succ m)
forall r rank.
(T (Shape rank) -> T r (T (Index rank))) -> Iterator r rank
Iterator ((T (Shape (Succ m)) -> T r (T (Index (Succ m))))
 -> Iterator r (Succ m))
-> (T (Shape (Succ m)) -> T r (T (Index (Succ m))))
-> Iterator r (Succ m)
forall a b. (a -> b) -> a -> b
$ (T (T ShapeTag m) -> T Int -> T r (T (Index (Succ m))))
-> T (Shape (Succ m)) -> T r (T (Index (Succ m)))
forall rank (val :: * -> *) tag a.
(Natural rank, Value val) =>
(val (T tag rank) -> val Int -> a) -> val (T tag (Succ rank)) -> a
switchR ((T (T ShapeTag m) -> T Int -> T r (T (Index (Succ m))))
 -> T (Shape (Succ m)) -> T r (T (Index (Succ m))))
-> (T (T ShapeTag m) -> T Int -> T r (T (Index (Succ m))))
-> T (Shape (Succ m))
-> T r (T (Index (Succ m)))
forall a b. (a -> b) -> a -> b
$ \T (T ShapeTag m)
sh T Int
n ->
       ((T (T IndexTag m), T Int) -> T (Index (Succ m)))
-> T r (T (T IndexTag m), T Int) -> T r (T (Index (Succ m)))
forall a b. (a -> b) -> T r a -> T r b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(T (T IndexTag m)
ix,T Int
i) -> T (T IndexTag m)
ixT (T IndexTag m) -> T Int -> T (Index (Succ m))
forall (val :: * -> *) tag rank.
Value val =>
val (T tag rank) -> val Int -> val (T tag (Succ rank))
#:.T Int
i) (T r (T (T IndexTag m), T Int) -> T r (T (Index (Succ m))))
-> T r (T (T IndexTag m), T Int) -> T r (T (Index (Succ m)))
forall a b. (a -> b) -> a -> b
$
       T r (T (T IndexTag m))
-> T r (T Int) -> T r (T (T IndexTag m), T Int)
forall a b r.
(Phi a, Phi b, Undefined a, Undefined b) =>
T r a -> T r b -> T r (a, b)
Iter.cartesian
         (T (T ShapeTag m) -> T r (T (T IndexTag m))
forall rank r.
Natural rank =>
T (Shape rank) -> T r (T (Index rank))
iterator T (T ShapeTag m)
sh)
         ((T Int -> CodeGenFunction r (T Bool)) -> T r (T Int) -> T r (T Int)
forall a r. (a -> CodeGenFunction r (T Bool)) -> T r a -> T r a
IterMV.takeWhile (CmpPredicate -> T Int -> T Int -> CodeGenFunction r (T Bool)
forall r.
CmpPredicate -> T Int -> T Int -> CodeGenFunction r (T Bool)
forall a r.
Comparison a =>
CmpPredicate -> T a -> T a -> CodeGenFunction r (T Bool)
MultiValue.cmp CmpPredicate
LLVM.CmpGT T Int
n) (T r (T Int) -> T r (T Int)) -> T r (T Int) -> T r (T Int)
forall a b. (a -> b) -> a -> b
$
          (T Int -> CodeGenFunction r (T Int)) -> T Int -> T r (T Int)
forall a r.
(Phi a, Undefined a) =>
(a -> CodeGenFunction r a) -> a -> T r a
Iter.iterate T Int -> CodeGenFunction r (T Int)
forall i r.
(Additive i, IntegerConstant i) =>
T i -> CodeGenFunction r (T i)
MultiValue.inc T Int
forall a. C a => T a
MultiValue.zero))


newtype Loop r state rank =
   Loop {
      forall r state rank.
Loop r state rank
-> (T (Index rank) -> state -> CodeGenFunction r state)
-> T (Shape rank)
-> state
-> CodeGenFunction r state
runLoop ::
         (MultiValue.T (Index rank) ->
          state ->
          LLVM.CodeGenFunction r state) ->
         MultiValue.T (Shape rank) ->
         state ->
         LLVM.CodeGenFunction r state
   }

loop ::
   (Unary.Natural rank, Tuple.Phi state) =>
   (MultiValue.T (Index rank) ->
    state ->
    LLVM.CodeGenFunction r state) ->
   MultiValue.T (Shape rank) ->
   state ->
   LLVM.CodeGenFunction r state
loop :: forall rank state r.
(Natural rank, Phi state) =>
(T (Index rank) -> state -> CodeGenFunction r state)
-> T (Shape rank) -> state -> CodeGenFunction r state
loop =
   Loop r state rank
-> (T (Index rank) -> state -> CodeGenFunction r state)
-> T (Shape rank)
-> state
-> CodeGenFunction r state
forall r state rank.
Loop r state rank
-> (T (Index rank) -> state -> CodeGenFunction r state)
-> T (Shape rank)
-> state
-> CodeGenFunction r state
runLoop (Loop r state rank
 -> (T (Index rank) -> state -> CodeGenFunction r state)
 -> T (Shape rank)
 -> state
 -> CodeGenFunction r state)
-> Loop r state rank
-> (T (Index rank) -> state -> CodeGenFunction r state)
-> T (Shape rank)
-> state
-> CodeGenFunction r state
forall a b. (a -> b) -> a -> b
$
   Loop r state Zero
-> (forall m. Natural m => Loop r state (Succ m))
-> Loop r state rank
forall n (f :: * -> *).
Natural n =>
f Zero -> (forall m. Natural m => f (Succ m)) -> f n
forall (f :: * -> *).
f Zero -> (forall m. Natural m => f (Succ m)) -> f rank
Unary.switchNat
      (((T (Index Zero) -> state -> CodeGenFunction r state)
 -> T (T ShapeTag Zero) -> state -> CodeGenFunction r state)
-> Loop r state Zero
forall r state rank.
((T (Index rank) -> state -> CodeGenFunction r state)
 -> T (Shape rank) -> state -> CodeGenFunction r state)
-> Loop r state rank
Loop (((T (Index Zero) -> state -> CodeGenFunction r state)
  -> T (T ShapeTag Zero) -> state -> CodeGenFunction r state)
 -> Loop r state Zero)
-> ((T (Index Zero) -> state -> CodeGenFunction r state)
    -> T (T ShapeTag Zero) -> state -> CodeGenFunction r state)
-> Loop r state Zero
forall a b. (a -> b) -> a -> b
$ \T (Index Zero) -> state -> CodeGenFunction r state
code T (T ShapeTag Zero)
_z -> T (Index Zero) -> state -> CodeGenFunction r state
code T (Index Zero)
forall (val :: * -> *) tag. Value val => val (T tag Zero)
z)
      (((T (Index (Succ m)) -> state -> CodeGenFunction r state)
 -> T (Shape (Succ m)) -> state -> CodeGenFunction r state)
-> Loop r state (Succ m)
forall r state rank.
((T (Index rank) -> state -> CodeGenFunction r state)
 -> T (Shape rank) -> state -> CodeGenFunction r state)
-> Loop r state rank
Loop (((T (Index (Succ m)) -> state -> CodeGenFunction r state)
  -> T (Shape (Succ m)) -> state -> CodeGenFunction r state)
 -> Loop r state (Succ m))
-> ((T (Index (Succ m)) -> state -> CodeGenFunction r state)
    -> T (Shape (Succ m)) -> state -> CodeGenFunction r state)
-> Loop r state (Succ m)
forall a b. (a -> b) -> a -> b
$ \T (Index (Succ m)) -> state -> CodeGenFunction r state
code -> (T (T ShapeTag m) -> T Int -> state -> CodeGenFunction r state)
-> T (Shape (Succ m)) -> state -> CodeGenFunction r state
forall rank (val :: * -> *) tag a.
(Natural rank, Value val) =>
(val (T tag rank) -> val Int -> a) -> val (T tag (Succ rank)) -> a
switchR ((T (T ShapeTag m) -> T Int -> state -> CodeGenFunction r state)
 -> T (Shape (Succ m)) -> state -> CodeGenFunction r state)
-> (T (T ShapeTag m) -> T Int -> state -> CodeGenFunction r state)
-> T (Shape (Succ m))
-> state
-> CodeGenFunction r state
forall a b. (a -> b) -> a -> b
$ \T (T ShapeTag m)
sh (MultiValue.Cons Repr Int
n) ->
         (T (Index m) -> state -> CodeGenFunction r state)
-> T (T ShapeTag m) -> state -> CodeGenFunction r state
forall rank state r.
(Natural rank, Phi state) =>
(T (Index rank) -> state -> CodeGenFunction r state)
-> T (Shape rank) -> state -> CodeGenFunction r state
loop
            (\T (Index m)
ix state
ptrStart ->
               ((state, Value Size) -> state)
-> CodeGenFunction r (state, Value Size) -> CodeGenFunction r state
forall a b. (a -> b) -> CodeGenFunction r a -> CodeGenFunction r b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (state, Value Size) -> state
forall a b. (a, b) -> a
fst (CodeGenFunction r (state, Value Size) -> CodeGenFunction r state)
-> CodeGenFunction r (state, Value Size) -> CodeGenFunction r state
forall a b. (a -> b) -> a -> b
$
               Value Size
-> (state, Value Size)
-> ((state, Value Size) -> CodeGenFunction r (state, Value Size))
-> CodeGenFunction r (state, Value Size)
forall s i r.
(Phi s, Num i, IsConst i, IsInteger i, CmpRet i, IsPrimitive i) =>
Value i -> s -> (s -> CodeGenFunction r s) -> CodeGenFunction r s
C.fixedLengthLoop Repr Int
Value Size
n (state
ptrStart, Value Size
forall a. Additive a => a
A.zero) (((state, Value Size) -> CodeGenFunction r (state, Value Size))
 -> CodeGenFunction r (state, Value Size))
-> ((state, Value Size) -> CodeGenFunction r (state, Value Size))
-> CodeGenFunction r (state, Value Size)
forall a b. (a -> b) -> a -> b
$ \(state
ptr, Value Size
k) ->
                  (state -> Value Size -> (state, Value Size))
-> CodeGenFunction r state
-> CodeGenFunction r (Value Size)
-> CodeGenFunction r (state, Value Size)
forall (m :: * -> *) a1 a2 r.
Monad m =>
(a1 -> a2 -> r) -> m a1 -> m a2 -> m r
liftM2 (,)
                     (T (Index (Succ m)) -> state -> CodeGenFunction r state
code (T (Index m)
ix T (Index m) -> T Int -> T (Index (Succ m))
forall (val :: * -> *) tag rank.
Value val =>
val (T tag rank) -> val Int -> val (T tag (Succ rank))
#:. Repr Int -> T Int
forall a. Repr a -> T a
MultiValue.Cons Repr Int
Value Size
k) state
ptr)
                     (Value Size -> CodeGenFunction r (Value Size)
forall a r.
(IsArithmetic a, IsConst a, Num a) =>
Value a -> CodeGenFunction r (Value a)
A.inc Value Size
k))
            T (T ShapeTag m)
sh)