{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE NoStarIsType #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneKindSignatures #-}
{-# LANGUAGE StrictData #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
module Data.Array.XArray where

import Control.DeepSeq (NFData)
import Data.Array.Internal qualified as OI
import Data.Array.Internal.RankedG qualified as ORG
import Data.Array.Internal.RankedS qualified as ORS
import Data.Array.Ranked qualified as ORB
import Data.Array.RankedS qualified as S
import Data.Coerce
import Data.Foldable (toList)
import Data.Kind
import Data.List.NonEmpty (NonEmpty)
import Data.Proxy
import Data.Type.Equality
import Data.Type.Ord
import Data.Vector.Storable qualified as VS
import Foreign.Storable (Storable)
import GHC.Generics (Generic)
import GHC.TypeLits

import Data.Array.Nested.Lemmas
import Data.Array.Nested.Mixed.Shape
import Data.Array.Nested.Permutation
import Data.Array.Nested.Types
import Data.Array.Strided.Orthotope


type XArray :: [Maybe Nat] -> Type -> Type
newtype XArray sh a = XArray (S.Array (Rank sh) a)
  deriving (Int -> XArray sh a -> ShowS
[XArray sh a] -> ShowS
XArray sh a -> String
(Int -> XArray sh a -> ShowS)
-> (XArray sh a -> String)
-> ([XArray sh a] -> ShowS)
-> Show (XArray sh a)
forall (sh :: [Maybe Nat]) a.
(Show a, Unbox a) =>
Int -> XArray sh a -> ShowS
forall (sh :: [Maybe Nat]) a.
(Show a, Unbox a) =>
[XArray sh a] -> ShowS
forall (sh :: [Maybe Nat]) a.
(Show a, Unbox a) =>
XArray sh a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall (sh :: [Maybe Nat]) a.
(Show a, Unbox a) =>
Int -> XArray sh a -> ShowS
showsPrec :: Int -> XArray sh a -> ShowS
$cshow :: forall (sh :: [Maybe Nat]) a.
(Show a, Unbox a) =>
XArray sh a -> String
show :: XArray sh a -> String
$cshowList :: forall (sh :: [Maybe Nat]) a.
(Show a, Unbox a) =>
[XArray sh a] -> ShowS
showList :: [XArray sh a] -> ShowS
Show, XArray sh a -> XArray sh a -> Bool
(XArray sh a -> XArray sh a -> Bool)
-> (XArray sh a -> XArray sh a -> Bool) -> Eq (XArray sh a)
forall (sh :: [Maybe Nat]) a.
(Eq a, Storable a) =>
XArray sh a -> XArray sh a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall (sh :: [Maybe Nat]) a.
(Eq a, Storable a) =>
XArray sh a -> XArray sh a -> Bool
== :: XArray sh a -> XArray sh a -> Bool
$c/= :: forall (sh :: [Maybe Nat]) a.
(Eq a, Storable a) =>
XArray sh a -> XArray sh a -> Bool
/= :: XArray sh a -> XArray sh a -> Bool
Eq, Eq (XArray sh a)
Eq (XArray sh a) =>
(XArray sh a -> XArray sh a -> Ordering)
-> (XArray sh a -> XArray sh a -> Bool)
-> (XArray sh a -> XArray sh a -> Bool)
-> (XArray sh a -> XArray sh a -> Bool)
-> (XArray sh a -> XArray sh a -> Bool)
-> (XArray sh a -> XArray sh a -> XArray sh a)
-> (XArray sh a -> XArray sh a -> XArray sh a)
-> Ord (XArray sh a)
XArray sh a -> XArray sh a -> Bool
XArray sh a -> XArray sh a -> Ordering
XArray sh a -> XArray sh a -> XArray sh a
forall (sh :: [Maybe Nat]) a.
(Ord a, Storable a) =>
Eq (XArray sh a)
forall (sh :: [Maybe Nat]) a.
(Ord a, Storable a) =>
XArray sh a -> XArray sh a -> Bool
forall (sh :: [Maybe Nat]) a.
(Ord a, Storable a) =>
XArray sh a -> XArray sh a -> Ordering
forall (sh :: [Maybe Nat]) a.
(Ord a, Storable a) =>
XArray sh a -> XArray sh a -> XArray sh a
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 :: forall (sh :: [Maybe Nat]) a.
(Ord a, Storable a) =>
XArray sh a -> XArray sh a -> Ordering
compare :: XArray sh a -> XArray sh a -> Ordering
$c< :: forall (sh :: [Maybe Nat]) a.
(Ord a, Storable a) =>
XArray sh a -> XArray sh a -> Bool
< :: XArray sh a -> XArray sh a -> Bool
$c<= :: forall (sh :: [Maybe Nat]) a.
(Ord a, Storable a) =>
XArray sh a -> XArray sh a -> Bool
<= :: XArray sh a -> XArray sh a -> Bool
$c> :: forall (sh :: [Maybe Nat]) a.
(Ord a, Storable a) =>
XArray sh a -> XArray sh a -> Bool
> :: XArray sh a -> XArray sh a -> Bool
$c>= :: forall (sh :: [Maybe Nat]) a.
(Ord a, Storable a) =>
XArray sh a -> XArray sh a -> Bool
>= :: XArray sh a -> XArray sh a -> Bool
$cmax :: forall (sh :: [Maybe Nat]) a.
(Ord a, Storable a) =>
XArray sh a -> XArray sh a -> XArray sh a
max :: XArray sh a -> XArray sh a -> XArray sh a
$cmin :: forall (sh :: [Maybe Nat]) a.
(Ord a, Storable a) =>
XArray sh a -> XArray sh a -> XArray sh a
min :: XArray sh a -> XArray sh a -> XArray sh a
Ord, (forall x. XArray sh a -> Rep (XArray sh a) x)
-> (forall x. Rep (XArray sh a) x -> XArray sh a)
-> Generic (XArray sh a)
forall (sh :: [Maybe Nat]) a x. Rep (XArray sh a) x -> XArray sh a
forall (sh :: [Maybe Nat]) a x. XArray sh a -> Rep (XArray sh a) x
forall x. Rep (XArray sh a) x -> XArray sh a
forall x. XArray sh a -> Rep (XArray sh a) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall (sh :: [Maybe Nat]) a x. XArray sh a -> Rep (XArray sh a) x
from :: forall x. XArray sh a -> Rep (XArray sh a) x
$cto :: forall (sh :: [Maybe Nat]) a x. Rep (XArray sh a) x -> XArray sh a
to :: forall x. Rep (XArray sh a) x -> XArray sh a
Generic)

instance NFData (XArray sh a)


shape :: forall sh a. StaticShX sh -> XArray sh a -> IShX sh
shape :: forall (sh :: [Maybe Nat]) a.
StaticShX sh -> XArray sh a -> IShX sh
shape = \StaticShX sh
ssh (XArray Array (Rank sh) a
arr) -> StaticShX sh -> [Int] -> IShX sh
forall (sh' :: [Maybe Nat]). StaticShX sh' -> [Int] -> IShX sh'
go StaticShX sh
ssh (Array (Rank sh) a -> [Int]
forall (n :: Nat) a. Array n a -> [Int]
S.shapeL Array (Rank sh) a
arr)
  where
    go :: StaticShX sh' -> [Int] -> IShX sh'
    go :: forall (sh' :: [Maybe Nat]). StaticShX sh' -> [Int] -> IShX sh'
go StaticShX sh'
ZKX [] = ShX sh' Int
forall (sh :: [Maybe Nat]) i. (sh ~ '[]) => ShX sh i
ZSX
    go (SMayNat () SNat n
n :!% StaticShX sh
ssh) (Int
i : [Int]
l) = ((n ~ 'Nothing) => () -> SMayNat Int SNat n)
-> (forall (m :: Nat).
    (n ~ 'Just m) =>
    SNat m -> SMayNat Int SNat n)
-> SMayNat () SNat n
-> SMayNat Int SNat n
forall {k} (n :: Maybe k) i r (f :: k -> Type).
((n ~ 'Nothing) => i -> r)
-> (forall (m :: k). (n ~ 'Just m) => f m -> r)
-> SMayNat i f n
-> r
fromSMayNat (\()
_ -> Int -> SMayNat Int SNat 'Nothing
forall {k} i (f :: k -> Type). i -> SMayNat i f 'Nothing
SUnknown Int
i) SNat m -> SMayNat Int SNat n
SNat m -> SMayNat Int SNat ('Just m)
forall (m :: Nat). (n ~ 'Just m) => SNat m -> SMayNat Int SNat n
forall {k} (f :: k -> Type) (n1 :: k) i.
f n1 -> SMayNat i f ('Just n1)
SKnown SMayNat () SNat n
n SMayNat Int SNat n -> ShX sh Int -> ShX sh' Int
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
SMayNat i SNat n -> ShX sh i -> ShX sh1 i
:$% StaticShX sh -> [Int] -> ShX sh Int
forall (sh' :: [Maybe Nat]). StaticShX sh' -> [Int] -> IShX sh'
go StaticShX sh
ssh [Int]
l
    go StaticShX sh'
_ [Int]
_ = String -> ShX sh' Int
forall a. HasCallStack => String -> a
error String
"Invalid shapeL"

fromVector :: forall sh a. Storable a => IShX sh -> VS.Vector a -> XArray sh a
fromVector :: forall (sh :: [Maybe Nat]) a.
Storable a =>
IShX sh -> Vector a -> XArray sh a
fromVector IShX sh
sh Vector a
v
  | Dict KnownNat (Rank sh)
Dict <- IShX sh -> Dict KnownNat (Rank sh)
forall (sh :: [Maybe Nat]) i. ShX sh i -> Dict KnownNat (Rank sh)
lemKnownNatRank IShX sh
sh
  = Array (Rank sh) a -> XArray sh a
forall (sh :: [Maybe Nat]) a. Array (Rank sh) a -> XArray sh a
XArray ([Int] -> Vector a -> Array (Rank sh) a
forall a (n :: Nat).
(Unbox a, KnownNat n) =>
[Int] -> Vector a -> Array n a
S.fromVector (IShX sh -> [Int]
forall (sh :: [Maybe Nat]). IShX sh -> [Int]
shxToList IShX sh
sh) Vector a
v)

toVector :: Storable a => XArray sh a -> VS.Vector a
toVector :: forall a (sh :: [Maybe Nat]). Storable a => XArray sh a -> Vector a
toVector (XArray Array (Rank sh) a
arr) = Array (Rank sh) a -> Vector a
forall (n :: Nat) a. Unbox a => Array n a -> Vector a
S.toVector Array (Rank sh) a
arr

-- | This allows observing the strides in the underlying orthotope array. This
-- can be useful for optimisation, but should be considered an implementation
-- detail: strides may change in new versions of this library without notice.
arrayStrides :: XArray sh a -> [Int]
arrayStrides :: forall (sh :: [Maybe Nat]) a. XArray sh a -> [Int]
arrayStrides (XArray (ORS.A (ORG.A [Int]
_ (OI.T [Int]
strides Int
_ Vector a
_)))) = [Int]
strides

scalar :: Storable a => a -> XArray '[] a
scalar :: forall a. Storable a => a -> XArray '[] a
scalar = Array 0 a -> XArray '[] a
Array (Rank '[]) a -> XArray '[] a
forall (sh :: [Maybe Nat]) a. Array (Rank sh) a -> XArray sh a
XArray (Array 0 a -> XArray '[] a)
-> (a -> Array 0 a) -> a -> XArray '[] a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Array 0 a
forall a. Unbox a => a -> Array 0 a
S.scalar

-- | Will throw if the array does not have the casted-to shape.
cast :: forall sh1 sh2 sh' a. Rank sh1 ~ Rank sh2
     => StaticShX sh1 -> IShX sh2 -> StaticShX sh'
     -> XArray (sh1 ++ sh') a -> XArray (sh2 ++ sh') a
cast :: forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat])
       (sh' :: [Maybe Nat]) a.
(Rank sh1 ~ Rank sh2) =>
StaticShX sh1
-> IShX sh2
-> StaticShX sh'
-> XArray (sh1 ++ sh') a
-> XArray (sh2 ++ sh') a
cast StaticShX sh1
ssh1 IShX sh2
sh2 StaticShX sh'
ssh' (XArray Array (Rank (sh1 ++ sh')) a
arr)
  | Rank (sh1 ++ sh') :~: (Rank sh1 + Rank sh')
Refl <- StaticShX sh1
-> StaticShX sh' -> Rank (sh1 ++ sh') :~: (Rank sh1 + Rank sh')
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]).
StaticShX sh1
-> StaticShX sh2 -> Rank (sh1 ++ sh2) :~: (Rank sh1 + Rank sh2)
lemRankApp StaticShX sh1
ssh1 StaticShX sh'
ssh'
  , Rank (sh2 ++ sh') :~: (Rank sh2 + Rank sh')
Refl <- StaticShX sh2
-> StaticShX sh' -> Rank (sh2 ++ sh') :~: (Rank sh2 + Rank sh')
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]).
StaticShX sh1
-> StaticShX sh2 -> Rank (sh1 ++ sh2) :~: (Rank sh1 + Rank sh2)
lemRankApp (IShX sh2 -> StaticShX sh2
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX IShX sh2
sh2) StaticShX sh'
ssh'
  = let arrsh :: IShX sh1
        (IShX sh1
arrsh, ShX sh' Int
_) = Proxy sh'
-> StaticShX sh1 -> ShX (sh1 ++ sh') Int -> (IShX sh1, ShX sh' Int)
forall (proxy :: [Maybe Nat] -> Type) (sh' :: [Maybe Nat])
       (sh :: [Maybe Nat]) i.
proxy sh'
-> StaticShX sh -> ShX (sh ++ sh') i -> (ShX sh i, ShX sh' i)
shxSplitApp (forall (t :: [Maybe Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @sh') StaticShX sh1
ssh1 (StaticShX (sh1 ++ sh')
-> XArray (sh1 ++ sh') a -> ShX (sh1 ++ sh') Int
forall (sh :: [Maybe Nat]) a.
StaticShX sh -> XArray sh a -> IShX sh
shape (StaticShX sh1 -> StaticShX sh' -> StaticShX (sh1 ++ sh')
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
StaticShX sh -> StaticShX sh' -> StaticShX (sh ++ sh')
ssxAppend StaticShX sh1
ssh1 StaticShX sh'
ssh') (Array (Rank (sh1 ++ sh')) a -> XArray (sh1 ++ sh') a
forall (sh :: [Maybe Nat]) a. Array (Rank sh) a -> XArray sh a
XArray Array (Rank (sh1 ++ sh')) a
arr))
    in if IShX sh1 -> [Int]
forall (sh :: [Maybe Nat]). IShX sh -> [Int]
shxToList IShX sh1
arrsh [Int] -> [Int] -> Bool
forall a. Eq a => a -> a -> Bool
== IShX sh2 -> [Int]
forall (sh :: [Maybe Nat]). IShX sh -> [Int]
shxToList IShX sh2
sh2
         then Array (Rank (sh2 ++ sh')) a -> XArray (sh2 ++ sh') a
forall (sh :: [Maybe Nat]) a. Array (Rank sh) a -> XArray sh a
XArray Array (Rank (sh1 ++ sh')) a
Array (Rank (sh2 ++ sh')) a
arr
         else String -> XArray (sh2 ++ sh') a
forall a. HasCallStack => String -> a
error (String -> XArray (sh2 ++ sh') a)
-> String -> XArray (sh2 ++ sh') a
forall a b. (a -> b) -> a -> b
$ String
"Data.Array.Mixed.cast: Cannot cast (" String -> ShowS
forall a. [a] -> [a] -> [a]
++ IShX sh1 -> String
forall a. Show a => a -> String
show IShX sh1
arrsh String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
") to (" String -> ShowS
forall a. [a] -> [a] -> [a]
++ IShX sh2 -> String
forall a. Show a => a -> String
show IShX sh2
sh2 String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"

unScalar :: Storable a => XArray '[] a -> a
unScalar :: forall a. Storable a => XArray '[] a -> a
unScalar (XArray Array (Rank '[]) a
a) = Array 0 a -> a
forall a. Unbox a => Array 0 a -> a
S.unScalar Array 0 a
Array (Rank '[]) a
a

replicate :: forall sh sh' a. Storable a => IShX sh -> StaticShX sh' -> XArray sh' a -> XArray (sh ++ sh') a
replicate :: forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) a.
Storable a =>
IShX sh -> StaticShX sh' -> XArray sh' a -> XArray (sh ++ sh') a
replicate IShX sh
sh StaticShX sh'
ssh' (XArray Array (Rank sh') a
arr)
  | Dict KnownNat (Rank sh')
Dict <- StaticShX sh' -> Dict KnownNat (Rank sh')
forall (sh :: [Maybe Nat]). StaticShX sh -> Dict KnownNat (Rank sh)
lemKnownNatRankSSX StaticShX sh'
ssh'
  , Dict KnownNat (Rank (sh ++ sh'))
Dict <- StaticShX (sh ++ sh') -> Dict KnownNat (Rank (sh ++ sh'))
forall (sh :: [Maybe Nat]). StaticShX sh -> Dict KnownNat (Rank sh)
lemKnownNatRankSSX (StaticShX sh -> StaticShX sh' -> StaticShX (sh ++ sh')
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
StaticShX sh -> StaticShX sh' -> StaticShX (sh ++ sh')
ssxAppend (IShX sh -> StaticShX sh
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX IShX sh
sh) StaticShX sh'
ssh')
  , Rank (sh ++ sh') :~: (Rank sh + Rank sh')
Refl <- StaticShX sh
-> StaticShX sh' -> Rank (sh ++ sh') :~: (Rank sh + Rank sh')
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]).
StaticShX sh1
-> StaticShX sh2 -> Rank (sh1 ++ sh2) :~: (Rank sh1 + Rank sh2)
lemRankApp (IShX sh -> StaticShX sh
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX IShX sh
sh) StaticShX sh'
ssh'
  = Array (Rank (sh ++ sh')) a -> XArray (sh ++ sh') a
forall (sh :: [Maybe Nat]) a. Array (Rank sh) a -> XArray sh a
XArray ([Int] -> Array (Rank (sh ++ sh')) a -> Array (Rank (sh ++ sh')) a
forall (n :: Nat) a. [Int] -> Array n a -> Array n a
S.stretch (IShX sh -> [Int]
forall (sh :: [Maybe Nat]). IShX sh -> [Int]
shxToList IShX sh
sh [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ Array (Rank sh') a -> [Int]
forall (n :: Nat) a. Array n a -> [Int]
S.shapeL Array (Rank sh') a
arr) (Array (Rank (sh ++ sh')) a -> Array (Rank (sh ++ sh')) a)
-> Array (Rank (sh ++ sh')) a -> Array (Rank (sh ++ sh')) a
forall a b. (a -> b) -> a -> b
$
            [Int] -> Array (Rank sh') a -> Array (Rank (sh ++ sh')) a
forall a (n :: Nat) (n' :: Nat).
(Unbox a, KnownNat n, KnownNat n') =>
[Int] -> Array n a -> Array n' a
S.reshape ((Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Int -> Int
forall a b. a -> b -> a
const Int
1) (IShX sh -> [Int]
forall (sh :: [Maybe Nat]). IShX sh -> [Int]
shxToList IShX sh
sh) [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ Array (Rank sh') a -> [Int]
forall (n :: Nat) a. Array n a -> [Int]
S.shapeL Array (Rank sh') a
arr)
              Array (Rank sh') a
arr)

replicateScal :: forall sh a. Storable a => IShX sh -> a -> XArray sh a
replicateScal :: forall (sh :: [Maybe Nat]) a.
Storable a =>
IShX sh -> a -> XArray sh a
replicateScal IShX sh
sh a
x
  | Dict KnownNat (Rank sh)
Dict <- IShX sh -> Dict KnownNat (Rank sh)
forall (sh :: [Maybe Nat]) i. ShX sh i -> Dict KnownNat (Rank sh)
lemKnownNatRank IShX sh
sh
  = Array (Rank sh) a -> XArray sh a
forall (sh :: [Maybe Nat]) a. Array (Rank sh) a -> XArray sh a
XArray ([Int] -> a -> Array (Rank sh) a
forall a (n :: Nat).
(Unbox a, KnownNat n) =>
[Int] -> a -> Array n a
S.constant (IShX sh -> [Int]
forall (sh :: [Maybe Nat]). IShX sh -> [Int]
shxToList IShX sh
sh) a
x)

generate :: Storable a => IShX sh -> (IIxX sh -> a) -> XArray sh a
generate :: forall a (sh :: [Maybe Nat]).
Storable a =>
IShX sh -> (IIxX sh -> a) -> XArray sh a
generate IShX sh
sh IIxX sh -> a
f = IShX sh -> Vector a -> XArray sh a
forall (sh :: [Maybe Nat]) a.
Storable a =>
IShX sh -> Vector a -> XArray sh a
fromVector IShX sh
sh (Vector a -> XArray sh a) -> Vector a -> XArray sh a
forall a b. (a -> b) -> a -> b
$ Int -> (Int -> a) -> Vector a
forall a. Storable a => Int -> (Int -> a) -> Vector a
VS.generate (IShX sh -> Int
forall (sh :: [Maybe Nat]). IShX sh -> Int
shxSize IShX sh
sh) (IIxX sh -> a
f (IIxX sh -> a) -> (Int -> IIxX sh) -> Int -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IShX sh -> Int -> IIxX sh
forall (sh :: [Maybe Nat]). IShX sh -> Int -> IIxX sh
ixxFromLinear IShX sh
sh)

-- generateM :: (Monad m, Storable a) => IShX sh -> (IIxX sh -> m a) -> m (XArray sh a)
-- generateM sh f | Dict <- lemKnownNatRank sh =
--   XArray . S.fromVector (shxShapeL sh)
--     <$> VS.generateM (shxSize sh) (f . ixxFromLinear sh)

indexPartial :: Storable a => XArray (sh ++ sh') a -> IIxX sh -> XArray sh' a
indexPartial :: forall a (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
Storable a =>
XArray (sh ++ sh') a -> IIxX sh -> XArray sh' a
indexPartial (XArray Array (Rank (sh ++ sh')) a
arr) IxX sh Int
ZIX = Array (Rank sh') a -> XArray sh' a
forall (sh :: [Maybe Nat]) a. Array (Rank sh) a -> XArray sh a
XArray Array (Rank sh') a
Array (Rank (sh ++ sh')) a
arr
indexPartial (XArray Array (Rank (sh ++ sh')) a
arr) (Int
i :.% IxX sh Int
idx) = XArray (sh ++ sh') a -> IxX sh Int -> XArray sh' a
forall a (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
Storable a =>
XArray (sh ++ sh') a -> IIxX sh -> XArray sh' a
indexPartial (Array (Rank (sh ++ sh')) a -> XArray (sh ++ sh') a
forall (sh :: [Maybe Nat]) a. Array (Rank sh) a -> XArray sh a
XArray (Array (1 + Rank (sh ++ sh')) a -> Int -> Array (Rank (sh ++ sh')) a
forall a (n :: Nat). Unbox a => Array (1 + n) a -> Int -> Array n a
S.index Array (1 + Rank (sh ++ sh')) a
Array (Rank (sh ++ sh')) a
arr Int
i)) IxX sh Int
idx

index :: forall sh a. Storable a => XArray sh a -> IIxX sh -> a
index :: forall (sh :: [Maybe Nat]) a.
Storable a =>
XArray sh a -> IIxX sh -> a
index XArray sh a
xarr IIxX sh
i
  | (sh ++ '[]) :~: sh
Refl <- forall (l :: [Maybe Nat]). (l ++ '[]) :~: l
forall {a} (l :: [a]). (l ++ '[]) :~: l
lemAppNil @sh
  = let XArray Array (Rank '[]) a
arr' = XArray (sh ++ '[]) a -> IIxX sh -> XArray '[] a
forall a (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
Storable a =>
XArray (sh ++ sh') a -> IIxX sh -> XArray sh' a
indexPartial XArray sh a
XArray (sh ++ '[]) a
xarr IIxX sh
i :: XArray '[] a
    in Array 0 a -> a
forall a. Unbox a => Array 0 a -> a
S.unScalar Array 0 a
Array (Rank '[]) a
arr'

append :: forall n m sh a. Storable a
       => StaticShX sh -> XArray (n : sh) a -> XArray (m : sh) a -> XArray (AddMaybe n m : sh) a
append :: forall (n :: Maybe Nat) (m :: Maybe Nat) (sh :: [Maybe Nat]) a.
Storable a =>
StaticShX sh
-> XArray (n : sh) a
-> XArray (m : sh) a
-> XArray (AddMaybe n m : sh) a
append StaticShX sh
ssh (XArray Array (Rank (n : sh)) a
a) (XArray Array (Rank (m : sh)) a
b)
  | Dict KnownNat (Rank sh)
Dict <- StaticShX sh -> Dict KnownNat (Rank sh)
forall (sh :: [Maybe Nat]). StaticShX sh -> Dict KnownNat (Rank sh)
lemKnownNatRankSSX StaticShX sh
ssh
  = Array (Rank (AddMaybe n m : sh)) a -> XArray (AddMaybe n m : sh) a
forall (sh :: [Maybe Nat]) a. Array (Rank sh) a -> XArray sh a
XArray (Array (Rank sh + 1) a
-> Array (Rank sh + 1) a -> Array (Rank sh + 1) a
forall a (n :: Nat).
(Unbox a, KnownNat n) =>
Array n a -> Array n a -> Array n a
S.append Array (Rank sh + 1) a
Array (Rank (n : sh)) a
a Array (Rank sh + 1) a
Array (Rank (m : sh)) a
b)

-- | All arrays must have the same shape, except possibly for the outermost
-- dimension.
concat :: Storable a
       => StaticShX sh -> NonEmpty (XArray (Nothing : sh) a) -> XArray (Nothing : sh) a
concat :: forall a (sh :: [Maybe Nat]).
Storable a =>
StaticShX sh
-> NonEmpty (XArray ('Nothing : sh) a) -> XArray ('Nothing : sh) a
concat StaticShX sh
ssh NonEmpty (XArray ('Nothing : sh) a)
l
  | Dict KnownNat (Rank sh)
Dict <- StaticShX sh -> Dict KnownNat (Rank sh)
forall (sh :: [Maybe Nat]). StaticShX sh -> Dict KnownNat (Rank sh)
lemKnownNatRankSSX StaticShX sh
ssh
  = Array (Rank ('Nothing : sh)) a -> XArray ('Nothing : sh) a
forall (sh :: [Maybe Nat]) a. Array (Rank sh) a -> XArray sh a
XArray ([Array (Rank sh + 1) a] -> Array (Rank sh + 1) a
forall a (n :: Nat).
(Unbox a, KnownNat n) =>
[Array n a] -> Array n a
S.concatOuter ([XArray ('Nothing : sh) a] -> [Array (Rank sh + 1) a]
forall a b. Coercible a b => a -> b
coerce (NonEmpty (XArray ('Nothing : sh) a) -> [XArray ('Nothing : sh) a]
forall a. NonEmpty a -> [a]
forall (t :: Type -> Type) a. Foldable t => t a -> [a]
toList NonEmpty (XArray ('Nothing : sh) a)
l)))

-- | If the prefix of the shape of the input array (@sh@) is empty (i.e.
-- contains a zero), then there is no way to deduce the full shape of the output
-- array (more precisely, the @sh2@ part): that could only come from calling
-- @f@, and there are no subarrays to call @f@ on. @orthotope@ errors out in
-- this case; we choose to fill the shape with zeros wherever we cannot deduce
-- what it should be.
--
-- For example, if:
--
-- @
-- arr :: XArray '[Just 3, Just 0, Just 4, Just 2, Nothing] Int   -- of shape [3, 0, 4, 2, 21]
-- f :: XArray '[Just 2, Nothing] Int -> XArray '[Just 5, Nothing, Just 17] Float
-- @
--
-- then:
--
-- @
-- rerank _ _ _ f arr :: XArray '[Just 3, Just 0, Just 4, Just 5, Nothing, Just 17] Float
-- @
--
-- and this result will have shape @[3, 0, 4, 5, 0, 17]@. Note the second @0@
-- in this shape: we don't know if @f@ intended to return an array with shape 0
-- here (it probably didn't), but there is no better number to put here absent
-- a subarray of the input to pass to @f@.
--
-- In this particular case the fact that @sh@ is empty was evident from the
-- type-level information, but the same situation occurs when @sh@ consists of
-- @Nothing@s, and some of those happen to be zero at runtime.
rerank :: forall sh sh1 sh2 a b.
          (Storable a, Storable b)
       => StaticShX sh -> StaticShX sh1 -> StaticShX sh2
       -> (XArray sh1 a -> XArray sh2 b)
       -> XArray (sh ++ sh1) a -> XArray (sh ++ sh2) b
rerank :: forall (sh :: [Maybe Nat]) (sh1 :: [Maybe Nat])
       (sh2 :: [Maybe Nat]) a b.
(Storable a, Storable b) =>
StaticShX sh
-> StaticShX sh1
-> StaticShX sh2
-> (XArray sh1 a -> XArray sh2 b)
-> XArray (sh ++ sh1) a
-> XArray (sh ++ sh2) b
rerank StaticShX sh
ssh StaticShX sh1
ssh1 StaticShX sh2
ssh2 XArray sh1 a -> XArray sh2 b
f xarr :: XArray (sh ++ sh1) a
xarr@(XArray Array (Rank (sh ++ sh1)) a
arr)
  | Dict KnownNat (Rank (sh ++ sh2))
Dict <- StaticShX (sh ++ sh2) -> Dict KnownNat (Rank (sh ++ sh2))
forall (sh :: [Maybe Nat]). StaticShX sh -> Dict KnownNat (Rank sh)
lemKnownNatRankSSX (StaticShX sh -> StaticShX sh2 -> StaticShX (sh ++ sh2)
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
StaticShX sh -> StaticShX sh' -> StaticShX (sh ++ sh')
ssxAppend StaticShX sh
ssh StaticShX sh2
ssh2)
  = let (ShX sh Int
sh, ShX sh1 Int
_) = Proxy sh1
-> StaticShX sh -> ShX (sh ++ sh1) Int -> (ShX sh Int, ShX sh1 Int)
forall (proxy :: [Maybe Nat] -> Type) (sh' :: [Maybe Nat])
       (sh :: [Maybe Nat]) i.
proxy sh'
-> StaticShX sh -> ShX (sh ++ sh') i -> (ShX sh i, ShX sh' i)
shxSplitApp (forall (t :: [Maybe Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @sh1) StaticShX sh
ssh (StaticShX (sh ++ sh1)
-> XArray (sh ++ sh1) a -> ShX (sh ++ sh1) Int
forall (sh :: [Maybe Nat]) a.
StaticShX sh -> XArray sh a -> IShX sh
shape (StaticShX sh -> StaticShX sh1 -> StaticShX (sh ++ sh1)
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
StaticShX sh -> StaticShX sh' -> StaticShX (sh ++ sh')
ssxAppend StaticShX sh
ssh StaticShX sh1
ssh1) XArray (sh ++ sh1) a
xarr)
    in if Int
0 Int -> [Int] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: Type -> Type) a.
(Foldable t, Eq a) =>
a -> t a -> Bool
`elem` ShX sh Int -> [Int]
forall (sh :: [Maybe Nat]). IShX sh -> [Int]
shxToList ShX sh Int
sh
         then Array (Rank (sh ++ sh2)) b -> XArray (sh ++ sh2) b
forall (sh :: [Maybe Nat]) a. Array (Rank sh) a -> XArray sh a
XArray ([Int] -> [b] -> Array (Rank (sh ++ sh2)) b
forall (n :: Nat) a.
(Unbox a, KnownNat n) =>
[Int] -> [a] -> Array n a
S.fromList (IShX (sh ++ sh2) -> [Int]
forall (sh :: [Maybe Nat]). IShX sh -> [Int]
shxToList (ShX sh Int -> ShX sh2 Int -> IShX (sh ++ sh2)
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) i.
ShX sh i -> ShX sh' i -> ShX (sh ++ sh') i
shxAppend ShX sh Int
sh (StaticShX sh2 -> ShX sh2 Int
forall (sh :: [Maybe Nat]). StaticShX sh -> IShX sh
shxCompleteZeros StaticShX sh2
ssh2))) [])
         else case () of
           () | Dict KnownNat (Rank sh)
Dict <- StaticShX sh -> Dict KnownNat (Rank sh)
forall (sh :: [Maybe Nat]). StaticShX sh -> Dict KnownNat (Rank sh)
lemKnownNatRankSSX StaticShX sh
ssh
              , Dict KnownNat (Rank sh2)
Dict <- StaticShX sh2 -> Dict KnownNat (Rank sh2)
forall (sh :: [Maybe Nat]). StaticShX sh -> Dict KnownNat (Rank sh)
lemKnownNatRankSSX StaticShX sh2
ssh2
              , Rank (sh ++ sh1) :~: (Rank sh + Rank sh1)
Refl <- StaticShX sh
-> StaticShX sh1 -> Rank (sh ++ sh1) :~: (Rank sh + Rank sh1)
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]).
StaticShX sh1
-> StaticShX sh2 -> Rank (sh1 ++ sh2) :~: (Rank sh1 + Rank sh2)
lemRankApp StaticShX sh
ssh StaticShX sh1
ssh1
              , Rank (sh ++ sh2) :~: (Rank sh + Rank sh2)
Refl <- StaticShX sh
-> StaticShX sh2 -> Rank (sh ++ sh2) :~: (Rank sh + Rank sh2)
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]).
StaticShX sh1
-> StaticShX sh2 -> Rank (sh1 ++ sh2) :~: (Rank sh1 + Rank sh2)
lemRankApp StaticShX sh
ssh StaticShX sh2
ssh2
              -> Array (Rank (sh ++ sh2)) b -> XArray (sh ++ sh2) b
forall (sh :: [Maybe Nat]) a. Array (Rank sh) a -> XArray sh a
XArray (forall (n :: Nat) (i :: Nat) (o :: Nat) a b.
(Unbox a, Unbox b, KnownNat n, KnownNat o, KnownNat (n + o),
 KnownNat (1 + o)) =>
(Array i a -> Array o b) -> Array (n + i) a -> Array (n + o) b
S.rerank @(Rank sh) @(Rank sh1) @(Rank sh2)
                          (\Array (Rank sh1) a
a -> let XArray Array (Rank sh2) b
r = XArray sh1 a -> XArray sh2 b
f (Array (Rank sh1) a -> XArray sh1 a
forall (sh :: [Maybe Nat]) a. Array (Rank sh) a -> XArray sh a
XArray Array (Rank sh1) a
a) in Array (Rank sh2) b
r)
                          Array (Rank sh + Rank sh1) a
Array (Rank (sh ++ sh1)) a
arr)

rerankTop :: forall sh1 sh2 sh a b.
             (Storable a, Storable b)
          => StaticShX sh1 -> StaticShX sh2 -> StaticShX sh
          -> (XArray sh1 a -> XArray sh2 b)
          -> XArray (sh1 ++ sh) a -> XArray (sh2 ++ sh) b
rerankTop :: forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat])
       (sh :: [Maybe Nat]) a b.
(Storable a, Storable b) =>
StaticShX sh1
-> StaticShX sh2
-> StaticShX sh
-> (XArray sh1 a -> XArray sh2 b)
-> XArray (sh1 ++ sh) a
-> XArray (sh2 ++ sh) b
rerankTop StaticShX sh1
ssh1 StaticShX sh2
ssh2 StaticShX sh
ssh XArray sh1 a -> XArray sh2 b
f = StaticShX sh
-> StaticShX sh2 -> XArray (sh ++ sh2) b -> XArray (sh2 ++ sh) b
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
StaticShX sh1
-> StaticShX sh2 -> XArray (sh1 ++ sh2) a -> XArray (sh2 ++ sh1) a
transpose2 StaticShX sh
ssh StaticShX sh2
ssh2 (XArray (sh ++ sh2) b -> XArray (sh2 ++ sh) b)
-> (XArray (sh1 ++ sh) a -> XArray (sh ++ sh2) b)
-> XArray (sh1 ++ sh) a
-> XArray (sh2 ++ sh) b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StaticShX sh
-> StaticShX sh1
-> StaticShX sh2
-> (XArray sh1 a -> XArray sh2 b)
-> XArray (sh ++ sh1) a
-> XArray (sh ++ sh2) b
forall (sh :: [Maybe Nat]) (sh1 :: [Maybe Nat])
       (sh2 :: [Maybe Nat]) a b.
(Storable a, Storable b) =>
StaticShX sh
-> StaticShX sh1
-> StaticShX sh2
-> (XArray sh1 a -> XArray sh2 b)
-> XArray (sh ++ sh1) a
-> XArray (sh ++ sh2) b
rerank StaticShX sh
ssh StaticShX sh1
ssh1 StaticShX sh2
ssh2 XArray sh1 a -> XArray sh2 b
f (XArray (sh ++ sh1) a -> XArray (sh ++ sh2) b)
-> (XArray (sh1 ++ sh) a -> XArray (sh ++ sh1) a)
-> XArray (sh1 ++ sh) a
-> XArray (sh ++ sh2) b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StaticShX sh1
-> StaticShX sh -> XArray (sh1 ++ sh) a -> XArray (sh ++ sh1) a
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
StaticShX sh1
-> StaticShX sh2 -> XArray (sh1 ++ sh2) a -> XArray (sh2 ++ sh1) a
transpose2 StaticShX sh1
ssh1 StaticShX sh
ssh

-- | The caveat about empty arrays at @rerank@ applies here too.
rerank2 :: forall sh sh1 sh2 a b c.
           (Storable a, Storable b, Storable c)
        => StaticShX sh -> StaticShX sh1 -> StaticShX sh2
        -> (XArray sh1 a -> XArray sh1 b -> XArray sh2 c)
        -> XArray (sh ++ sh1) a -> XArray (sh ++ sh1) b -> XArray (sh ++ sh2) c
rerank2 :: forall (sh :: [Maybe Nat]) (sh1 :: [Maybe Nat])
       (sh2 :: [Maybe Nat]) a b c.
(Storable a, Storable b, Storable c) =>
StaticShX sh
-> StaticShX sh1
-> StaticShX sh2
-> (XArray sh1 a -> XArray sh1 b -> XArray sh2 c)
-> XArray (sh ++ sh1) a
-> XArray (sh ++ sh1) b
-> XArray (sh ++ sh2) c
rerank2 StaticShX sh
ssh StaticShX sh1
ssh1 StaticShX sh2
ssh2 XArray sh1 a -> XArray sh1 b -> XArray sh2 c
f xarr1 :: XArray (sh ++ sh1) a
xarr1@(XArray Array (Rank (sh ++ sh1)) a
arr1) (XArray Array (Rank (sh ++ sh1)) b
arr2)
  | Dict KnownNat (Rank (sh ++ sh2))
Dict <- StaticShX (sh ++ sh2) -> Dict KnownNat (Rank (sh ++ sh2))
forall (sh :: [Maybe Nat]). StaticShX sh -> Dict KnownNat (Rank sh)
lemKnownNatRankSSX (StaticShX sh -> StaticShX sh2 -> StaticShX (sh ++ sh2)
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
StaticShX sh -> StaticShX sh' -> StaticShX (sh ++ sh')
ssxAppend StaticShX sh
ssh StaticShX sh2
ssh2)
  = let (ShX sh Int
sh, ShX sh1 Int
_) = Proxy sh1
-> StaticShX sh -> ShX (sh ++ sh1) Int -> (ShX sh Int, ShX sh1 Int)
forall (proxy :: [Maybe Nat] -> Type) (sh' :: [Maybe Nat])
       (sh :: [Maybe Nat]) i.
proxy sh'
-> StaticShX sh -> ShX (sh ++ sh') i -> (ShX sh i, ShX sh' i)
shxSplitApp (forall (t :: [Maybe Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @sh1) StaticShX sh
ssh (StaticShX (sh ++ sh1)
-> XArray (sh ++ sh1) a -> ShX (sh ++ sh1) Int
forall (sh :: [Maybe Nat]) a.
StaticShX sh -> XArray sh a -> IShX sh
shape (StaticShX sh -> StaticShX sh1 -> StaticShX (sh ++ sh1)
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
StaticShX sh -> StaticShX sh' -> StaticShX (sh ++ sh')
ssxAppend StaticShX sh
ssh StaticShX sh1
ssh1) XArray (sh ++ sh1) a
xarr1)
    in if Int
0 Int -> [Int] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: Type -> Type) a.
(Foldable t, Eq a) =>
a -> t a -> Bool
`elem` ShX sh Int -> [Int]
forall (sh :: [Maybe Nat]). IShX sh -> [Int]
shxToList ShX sh Int
sh
         then Array (Rank (sh ++ sh2)) c -> XArray (sh ++ sh2) c
forall (sh :: [Maybe Nat]) a. Array (Rank sh) a -> XArray sh a
XArray ([Int] -> [c] -> Array (Rank (sh ++ sh2)) c
forall (n :: Nat) a.
(Unbox a, KnownNat n) =>
[Int] -> [a] -> Array n a
S.fromList (IShX (sh ++ sh2) -> [Int]
forall (sh :: [Maybe Nat]). IShX sh -> [Int]
shxToList (ShX sh Int -> ShX sh2 Int -> IShX (sh ++ sh2)
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) i.
ShX sh i -> ShX sh' i -> ShX (sh ++ sh') i
shxAppend ShX sh Int
sh (StaticShX sh2 -> ShX sh2 Int
forall (sh :: [Maybe Nat]). StaticShX sh -> IShX sh
shxCompleteZeros StaticShX sh2
ssh2))) [])
         else case () of
           () | Dict KnownNat (Rank sh)
Dict <- StaticShX sh -> Dict KnownNat (Rank sh)
forall (sh :: [Maybe Nat]). StaticShX sh -> Dict KnownNat (Rank sh)
lemKnownNatRankSSX StaticShX sh
ssh
              , Dict KnownNat (Rank sh2)
Dict <- StaticShX sh2 -> Dict KnownNat (Rank sh2)
forall (sh :: [Maybe Nat]). StaticShX sh -> Dict KnownNat (Rank sh)
lemKnownNatRankSSX StaticShX sh2
ssh2
              , Rank (sh ++ sh1) :~: (Rank sh + Rank sh1)
Refl <- StaticShX sh
-> StaticShX sh1 -> Rank (sh ++ sh1) :~: (Rank sh + Rank sh1)
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]).
StaticShX sh1
-> StaticShX sh2 -> Rank (sh1 ++ sh2) :~: (Rank sh1 + Rank sh2)
lemRankApp StaticShX sh
ssh StaticShX sh1
ssh1
              , Rank (sh ++ sh2) :~: (Rank sh + Rank sh2)
Refl <- StaticShX sh
-> StaticShX sh2 -> Rank (sh ++ sh2) :~: (Rank sh + Rank sh2)
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]).
StaticShX sh1
-> StaticShX sh2 -> Rank (sh1 ++ sh2) :~: (Rank sh1 + Rank sh2)
lemRankApp StaticShX sh
ssh StaticShX sh2
ssh2
              -> Array (Rank (sh ++ sh2)) c -> XArray (sh ++ sh2) c
forall (sh :: [Maybe Nat]) a. Array (Rank sh) a -> XArray sh a
XArray (forall (n :: Nat) (i :: Nat) (o :: Nat) a b c.
(Unbox a, Unbox b, Unbox c, KnownNat n, KnownNat o,
 KnownNat (n + o), KnownNat (1 + o)) =>
(Array i a -> Array i b -> Array o c)
-> Array (n + i) a -> Array (n + i) b -> Array (n + o) c
S.rerank2 @(Rank sh) @(Rank sh1) @(Rank sh2)
                          (\Array (Rank sh1) a
a Array (Rank sh1) b
b -> let XArray Array (Rank sh2) c
r = XArray sh1 a -> XArray sh1 b -> XArray sh2 c
f (Array (Rank sh1) a -> XArray sh1 a
forall (sh :: [Maybe Nat]) a. Array (Rank sh) a -> XArray sh a
XArray Array (Rank sh1) a
a) (Array (Rank sh1) b -> XArray sh1 b
forall (sh :: [Maybe Nat]) a. Array (Rank sh) a -> XArray sh a
XArray Array (Rank sh1) b
b) in Array (Rank sh2) c
r)
                          Array (Rank sh + Rank sh1) a
Array (Rank (sh ++ sh1)) a
arr1 Array (Rank sh + Rank sh1) b
Array (Rank (sh ++ sh1)) b
arr2)

-- | The list argument gives indices into the original dimension list.
transpose :: forall is sh a. (IsPermutation is, Rank is <= Rank sh)
          => StaticShX sh
          -> Perm is
          -> XArray sh a
          -> XArray (PermutePrefix is sh) a
transpose :: forall (is :: [Nat]) (sh :: [Maybe Nat]) a.
(IsPermutation is, Rank is <= Rank sh) =>
StaticShX sh
-> Perm is -> XArray sh a -> XArray (PermutePrefix is sh) a
transpose StaticShX sh
ssh Perm is
perm (XArray Array (Rank sh) a
arr)
  | Dict KnownNat (Rank sh)
Dict <- StaticShX sh -> Dict KnownNat (Rank sh)
forall (sh :: [Maybe Nat]). StaticShX sh -> Dict KnownNat (Rank sh)
lemKnownNatRankSSX StaticShX sh
ssh
  , Rank (Permute is (TakeLen is sh) ++ DropLen is sh)
:~: (Rank (Permute is (TakeLen is sh)) + Rank (DropLen is sh))
Refl <- StaticShX (Permute is (TakeLen is sh))
-> StaticShX (DropLen is sh)
-> Rank (Permute is (TakeLen is sh) ++ DropLen is sh)
   :~: (Rank (Permute is (TakeLen is sh)) + Rank (DropLen is sh))
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]).
StaticShX sh1
-> StaticShX sh2 -> Rank (sh1 ++ sh2) :~: (Rank sh1 + Rank sh2)
lemRankApp (Perm is
-> StaticShX (TakeLen is sh)
-> StaticShX (Permute is (TakeLen is sh))
forall (is :: [Nat]) (sh :: [Maybe Nat]).
Perm is -> StaticShX sh -> StaticShX (Permute is sh)
ssxPermute Perm is
perm (Perm is -> StaticShX sh -> StaticShX (TakeLen is sh)
forall (is :: [Nat]) (sh :: [Maybe Nat]).
Perm is -> StaticShX sh -> StaticShX (TakeLen is sh)
ssxTakeLen Perm is
perm StaticShX sh
ssh)) (Perm is -> StaticShX sh -> StaticShX (DropLen is sh)
forall (is :: [Nat]) (sh :: [Maybe Nat]).
Perm is -> StaticShX sh -> StaticShX (DropLen is sh)
ssxDropLen Perm is
perm StaticShX sh
ssh)
  , Rank (Permute is (TakeLen is sh)) :~: Rank is
Refl <- Proxy (TakeLen is sh)
-> Perm is -> Rank (Permute is (TakeLen is sh)) :~: Rank is
forall {a} (sh :: [a]) (is :: [Nat]).
Proxy sh -> Perm is -> Rank (Permute is sh) :~: Rank is
lemRankPermute (forall (t :: [Maybe Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @(TakeLen is sh)) Perm is
perm
  , Rank (DropLen is sh) :~: (Rank sh - Rank is)
Refl <- StaticShX sh
-> Perm is -> Rank (DropLen is sh) :~: (Rank sh - Rank is)
forall (is :: [Nat]) (sh :: [Maybe Nat]).
(Rank is <= Rank sh) =>
StaticShX sh
-> Perm is -> Rank (DropLen is sh) :~: (Rank sh - Rank is)
lemRankDropLen StaticShX sh
ssh Perm is
perm
  = Array (Rank (Permute is (TakeLen is sh) ++ DropLen is sh)) a
-> XArray (Permute is (TakeLen is sh) ++ DropLen is sh) a
forall (sh :: [Maybe Nat]) a. Array (Rank sh) a -> XArray sh a
XArray ([Int]
-> Array (Rank is + (Rank sh - Rank is)) a
-> Array (Rank is + (Rank sh - Rank is)) a
forall (n :: Nat) a. KnownNat n => [Int] -> Array n a -> Array n a
S.transpose (Perm is -> [Int]
forall (list :: [Nat]). Perm list -> [Int]
permToList' Perm is
perm) Array (Rank is + (Rank sh - Rank is)) a
Array (Rank sh) a
arr)

-- | The list argument gives indices into the original dimension list.
--
-- The permutation (the list) must have length <= @n@. If it is longer, this
-- function throws.
transposeUntyped :: forall n sh a.
                    SNat n -> StaticShX sh -> [Int]
                 -> XArray (Replicate n Nothing ++ sh) a -> XArray (Replicate n Nothing ++ sh) a
transposeUntyped :: forall (n :: Nat) (sh :: [Maybe Nat]) a.
SNat n
-> StaticShX sh
-> [Int]
-> XArray (Replicate n 'Nothing ++ sh) a
-> XArray (Replicate n 'Nothing ++ sh) a
transposeUntyped SNat n
sn StaticShX sh
ssh [Int]
perm (XArray Array (Rank (Replicate n 'Nothing ++ sh)) a
arr)
  | [Int] -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [Int]
perm Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= SNat n -> Int
forall (n :: Nat). SNat n -> Int
fromSNat' SNat n
sn
  , Dict KnownNat (Rank (Replicate n 'Nothing ++ sh))
Dict <- StaticShX (Replicate n 'Nothing ++ sh)
-> Dict KnownNat (Rank (Replicate n 'Nothing ++ sh))
forall (sh :: [Maybe Nat]). StaticShX sh -> Dict KnownNat (Rank sh)
lemKnownNatRankSSX (StaticShX (Replicate n 'Nothing)
-> StaticShX sh -> StaticShX (Replicate n 'Nothing ++ sh)
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
StaticShX sh -> StaticShX sh' -> StaticShX (sh ++ sh')
ssxAppend (SNat n -> StaticShX (Replicate n 'Nothing)
forall (n :: Nat). SNat n -> StaticShX (Replicate n 'Nothing)
ssxReplicate SNat n
sn) StaticShX sh
ssh)
  = Array (Rank (Replicate n 'Nothing ++ sh)) a
-> XArray (Replicate n 'Nothing ++ sh) a
forall (sh :: [Maybe Nat]) a. Array (Rank sh) a -> XArray sh a
XArray ([Int]
-> Array (Rank (Replicate n 'Nothing ++ sh)) a
-> Array (Rank (Replicate n 'Nothing ++ sh)) a
forall (n :: Nat) a. KnownNat n => [Int] -> Array n a -> Array n a
S.transpose [Int]
perm Array (Rank (Replicate n 'Nothing ++ sh)) a
arr)
  | Bool
otherwise
  = String -> XArray (Replicate n 'Nothing ++ sh) a
forall a. HasCallStack => String -> a
error String
"Data.Array.Mixed.transposeUntyped: Permutation longer than length of unshaped prefix of shape type"

transpose2 :: forall sh1 sh2 a.
              StaticShX sh1 -> StaticShX sh2
           -> XArray (sh1 ++ sh2) a -> XArray (sh2 ++ sh1) a
transpose2 :: forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
StaticShX sh1
-> StaticShX sh2 -> XArray (sh1 ++ sh2) a -> XArray (sh2 ++ sh1) a
transpose2 StaticShX sh1
ssh1 StaticShX sh2
ssh2 (XArray Array (Rank (sh1 ++ sh2)) a
arr)
  | Rank (sh1 ++ sh2) :~: (Rank sh1 + Rank sh2)
Refl <- StaticShX sh1
-> StaticShX sh2 -> Rank (sh1 ++ sh2) :~: (Rank sh1 + Rank sh2)
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]).
StaticShX sh1
-> StaticShX sh2 -> Rank (sh1 ++ sh2) :~: (Rank sh1 + Rank sh2)
lemRankApp StaticShX sh1
ssh1 StaticShX sh2
ssh2
  , Rank (sh2 ++ sh1) :~: (Rank sh2 + Rank sh1)
Refl <- StaticShX sh2
-> StaticShX sh1 -> Rank (sh2 ++ sh1) :~: (Rank sh2 + Rank sh1)
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]).
StaticShX sh1
-> StaticShX sh2 -> Rank (sh1 ++ sh2) :~: (Rank sh1 + Rank sh2)
lemRankApp StaticShX sh2
ssh2 StaticShX sh1
ssh1
  , Dict KnownNat (Rank (sh1 ++ sh2))
Dict <- StaticShX (sh1 ++ sh2) -> Dict KnownNat (Rank (sh1 ++ sh2))
forall (sh :: [Maybe Nat]). StaticShX sh -> Dict KnownNat (Rank sh)
lemKnownNatRankSSX (StaticShX sh1 -> StaticShX sh2 -> StaticShX (sh1 ++ sh2)
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
StaticShX sh -> StaticShX sh' -> StaticShX (sh ++ sh')
ssxAppend StaticShX sh1
ssh1 StaticShX sh2
ssh2)
  , Dict KnownNat (Rank (sh2 ++ sh1))
Dict <- StaticShX (sh2 ++ sh1) -> Dict KnownNat (Rank (sh2 ++ sh1))
forall (sh :: [Maybe Nat]). StaticShX sh -> Dict KnownNat (Rank sh)
lemKnownNatRankSSX (StaticShX sh2 -> StaticShX sh1 -> StaticShX (sh2 ++ sh1)
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
StaticShX sh -> StaticShX sh' -> StaticShX (sh ++ sh')
ssxAppend StaticShX sh2
ssh2 StaticShX sh1
ssh1)
  , Rank (sh1 ++ sh2) :~: Rank (sh2 ++ sh1)
Refl <- StaticShX sh1
-> StaticShX sh2 -> Rank (sh1 ++ sh2) :~: Rank (sh2 ++ sh1)
forall {a} (proxy :: [a] -> Type) (sh1 :: [a]) (sh2 :: [a]).
proxy sh1 -> proxy sh2 -> Rank (sh1 ++ sh2) :~: Rank (sh2 ++ sh1)
lemRankAppComm StaticShX sh1
ssh1 StaticShX sh2
ssh2
  , let n1 :: Int
n1 = StaticShX sh1 -> Int
forall (sh :: [Maybe Nat]). StaticShX sh -> Int
ssxLength StaticShX sh1
ssh1
  = Array (Rank (sh2 ++ sh1)) a -> XArray (sh2 ++ sh1) a
forall (sh :: [Maybe Nat]) a. Array (Rank sh) a -> XArray sh a
XArray ([Int] -> Array (Rank (sh1 ++ sh2)) a -> Array (Rank (sh1 ++ sh2)) a
forall (n :: Nat) a. KnownNat n => [Int] -> Array n a -> Array n a
S.transpose (StaticShX sh2 -> Int -> [Int]
forall (sh :: [Maybe Nat]). StaticShX sh -> Int -> [Int]
ssxIotaFrom StaticShX sh2
ssh2 Int
n1 [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ StaticShX sh1 -> Int -> [Int]
forall (sh :: [Maybe Nat]). StaticShX sh -> Int -> [Int]
ssxIotaFrom StaticShX sh1
ssh1 Int
0) Array (Rank (sh1 ++ sh2)) a
arr)

sumFull :: (Storable a, NumElt a) => StaticShX sh -> XArray sh a -> a
sumFull :: forall a (sh :: [Maybe Nat]).
(Storable a, NumElt a) =>
StaticShX sh -> XArray sh a -> a
sumFull StaticShX sh
_ (XArray Array (Rank sh) a
arr) =
  Array 0 a -> a
forall a. Unbox a => Array 0 a -> a
S.unScalar (Array 0 a -> a) -> Array 0 a -> a
forall a b. (a -> b) -> a -> b
$
    (Array 1 a -> Array 0 a) -> Array 1 a -> Array 0 a
forall (n :: Nat) a (n' :: Nat) b.
(Array n a -> Array n' b) -> Array n a -> Array n' b
liftO1 (SNat 0 -> Array (0 + 1) a -> Array 0 a
forall (n :: Nat). SNat n -> Array (n + 1) a -> Array n a
forall a (n :: Nat).
NumElt a =>
SNat n -> Array (n + 1) a -> Array n a
numEltSum1Inner (forall (n :: Nat). KnownNat n => SNat n
SNat @0)) (Array 1 a -> Array 0 a) -> Array 1 a -> Array 0 a
forall a b. (a -> b) -> a -> b
$
      [Int] -> Vector a -> Array 1 a
forall a (n :: Nat).
(Unbox a, KnownNat n) =>
[Int] -> Vector a -> Array n a
S.fromVector [[Int] -> Int
forall a. Num a => [a] -> a
forall (t :: Type -> Type) a. (Foldable t, Num a) => t a -> a
product (Array (Rank sh) a -> [Int]
forall (n :: Nat) a. Array n a -> [Int]
S.shapeL Array (Rank sh) a
arr)] (Vector a -> Array 1 a) -> Vector a -> Array 1 a
forall a b. (a -> b) -> a -> b
$
        Array (Rank sh) a -> Vector a
forall (n :: Nat) a. Unbox a => Array n a -> Vector a
S.toVector Array (Rank sh) a
arr

sumInner :: forall sh sh' a. (Storable a, NumElt a)
         => StaticShX sh -> StaticShX sh' -> XArray (sh ++ sh') a -> XArray sh a
sumInner :: forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) a.
(Storable a, NumElt a) =>
StaticShX sh
-> StaticShX sh' -> XArray (sh ++ sh') a -> XArray sh a
sumInner StaticShX sh
ssh StaticShX sh'
ssh' XArray (sh ++ sh') a
arr
  | (sh ++ '[]) :~: sh
Refl <- forall (l :: [Maybe Nat]). (l ++ '[]) :~: l
forall {a} (l :: [a]). (l ++ '[]) :~: l
lemAppNil @sh
  = let (ShX sh Int
_, ShX sh' Int
sh') = Proxy sh'
-> StaticShX sh -> ShX (sh ++ sh') Int -> (ShX sh Int, ShX sh' Int)
forall (proxy :: [Maybe Nat] -> Type) (sh' :: [Maybe Nat])
       (sh :: [Maybe Nat]) i.
proxy sh'
-> StaticShX sh -> ShX (sh ++ sh') i -> (ShX sh i, ShX sh' i)
shxSplitApp (forall (t :: [Maybe Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @sh') StaticShX sh
ssh (StaticShX (sh ++ sh')
-> XArray (sh ++ sh') a -> ShX (sh ++ sh') Int
forall (sh :: [Maybe Nat]) a.
StaticShX sh -> XArray sh a -> IShX sh
shape (StaticShX sh -> StaticShX sh' -> StaticShX (sh ++ sh')
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
StaticShX sh -> StaticShX sh' -> StaticShX (sh ++ sh')
ssxAppend StaticShX sh
ssh StaticShX sh'
ssh') XArray (sh ++ sh') a
arr)
        sh'F :: ShX '[Flatten' 1 sh'] Int
sh'F = ShX sh' Int -> SMayNat Int SNat (Flatten' 1 sh')
forall (sh :: [Maybe Nat]).
IShX sh -> SMayNat Int SNat (Flatten sh)
shxFlatten ShX sh' Int
sh' SMayNat Int SNat (Flatten' 1 sh')
-> ShX '[] Int -> ShX '[Flatten' 1 sh'] Int
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
SMayNat i SNat n -> ShX sh i -> ShX sh1 i
:$% ShX '[] Int
forall (sh :: [Maybe Nat]) i. (sh ~ '[]) => ShX sh i
ZSX
        ssh'F :: StaticShX '[Flatten' 1 sh']
ssh'F = ShX '[Flatten' 1 sh'] Int -> StaticShX '[Flatten' 1 sh']
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX ShX '[Flatten' 1 sh'] Int
sh'F

        go :: XArray (sh ++ '[Flatten sh']) a -> XArray sh a
        go :: XArray (sh ++ '[Flatten' 1 sh']) a -> XArray sh a
go (XArray Array (Rank (sh ++ '[Flatten' 1 sh'])) a
arr')
          | Rank (sh ++ '[Flatten' 1 sh'])
:~: (Rank sh + Rank '[Flatten' 1 sh'])
Refl <- StaticShX sh
-> StaticShX '[Flatten' 1 sh']
-> Rank (sh ++ '[Flatten' 1 sh'])
   :~: (Rank sh + Rank '[Flatten' 1 sh'])
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]).
StaticShX sh1
-> StaticShX sh2 -> Rank (sh1 ++ sh2) :~: (Rank sh1 + Rank sh2)
lemRankApp StaticShX sh
ssh StaticShX '[Flatten' 1 sh']
ssh'F
          , let sn :: SNat (Rank sh)
sn = ListX sh (SMayNat () SNat) -> SNat (Rank sh)
forall (sh :: [Maybe Nat]) (f :: Maybe Nat -> Type).
ListX sh f -> SNat (Rank sh)
listxRank (let StaticShX ListX sh (SMayNat () SNat)
l = StaticShX sh
ssh in ListX sh (SMayNat () SNat)
l)
          = Array (Rank sh) a -> XArray sh a
forall (sh :: [Maybe Nat]) a. Array (Rank sh) a -> XArray sh a
XArray ((Array (Rank sh + 1) a -> Array (Rank sh) a)
-> Array (Rank sh + 1) a -> Array (Rank sh) a
forall (n :: Nat) a (n' :: Nat) b.
(Array n a -> Array n' b) -> Array n a -> Array n' b
liftO1 (SNat (Rank sh) -> Array (Rank sh + 1) a -> Array (Rank sh) a
forall (n :: Nat). SNat n -> Array (n + 1) a -> Array n a
forall a (n :: Nat).
NumElt a =>
SNat n -> Array (n + 1) a -> Array n a
numEltSum1Inner SNat (Rank sh)
sn) Array (Rank sh + 1) a
Array (Rank (sh ++ '[Flatten' 1 sh'])) a
arr')

    in XArray (sh ++ '[Flatten' 1 sh']) a -> XArray sh a
go (XArray (sh ++ '[Flatten' 1 sh']) a -> XArray sh a)
-> XArray (sh ++ '[Flatten' 1 sh']) a -> XArray sh a
forall a b. (a -> b) -> a -> b
$
       StaticShX '[Flatten' 1 sh']
-> StaticShX sh
-> XArray ('[Flatten' 1 sh'] ++ sh) a
-> XArray (sh ++ '[Flatten' 1 sh']) a
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
StaticShX sh1
-> StaticShX sh2 -> XArray (sh1 ++ sh2) a -> XArray (sh2 ++ sh1) a
transpose2 StaticShX '[Flatten' 1 sh']
ssh'F StaticShX sh
ssh (XArray ('[Flatten' 1 sh'] ++ sh) a
 -> XArray (sh ++ '[Flatten' 1 sh']) a)
-> XArray ('[Flatten' 1 sh'] ++ sh) a
-> XArray (sh ++ '[Flatten' 1 sh']) a
forall a b. (a -> b) -> a -> b
$
       StaticShX sh'
-> StaticShX sh
-> ShX '[Flatten' 1 sh'] Int
-> XArray (sh' ++ sh) a
-> XArray ('[Flatten' 1 sh'] ++ sh) a
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat])
       (sh' :: [Maybe Nat]) a.
Storable a =>
StaticShX sh1
-> StaticShX sh'
-> IShX sh2
-> XArray (sh1 ++ sh') a
-> XArray (sh2 ++ sh') a
reshapePartial StaticShX sh'
ssh' StaticShX sh
ssh ShX '[Flatten' 1 sh'] Int
sh'F (XArray (sh' ++ sh) a -> XArray ('[Flatten' 1 sh'] ++ sh) a)
-> XArray (sh' ++ sh) a -> XArray ('[Flatten' 1 sh'] ++ sh) a
forall a b. (a -> b) -> a -> b
$
       StaticShX sh
-> StaticShX sh' -> XArray (sh ++ sh') a -> XArray (sh' ++ sh) a
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
StaticShX sh1
-> StaticShX sh2 -> XArray (sh1 ++ sh2) a -> XArray (sh2 ++ sh1) a
transpose2 StaticShX sh
ssh StaticShX sh'
ssh' (XArray (sh ++ sh') a -> XArray (sh' ++ sh) a)
-> XArray (sh ++ sh') a -> XArray (sh' ++ sh) a
forall a b. (a -> b) -> a -> b
$
         XArray (sh ++ sh') a
arr

sumOuter :: forall sh sh' a. (Storable a, NumElt a)
         => StaticShX sh -> StaticShX sh' -> XArray (sh ++ sh') a -> XArray sh' a
sumOuter :: forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) a.
(Storable a, NumElt a) =>
StaticShX sh
-> StaticShX sh' -> XArray (sh ++ sh') a -> XArray sh' a
sumOuter StaticShX sh
ssh StaticShX sh'
ssh' XArray (sh ++ sh') a
arr
  | (sh ++ '[]) :~: sh
Refl <- forall (l :: [Maybe Nat]). (l ++ '[]) :~: l
forall {a} (l :: [a]). (l ++ '[]) :~: l
lemAppNil @sh
  = let (ShX sh Int
sh, ShX sh' Int
_) = Proxy sh'
-> StaticShX sh -> ShX (sh ++ sh') Int -> (ShX sh Int, ShX sh' Int)
forall (proxy :: [Maybe Nat] -> Type) (sh' :: [Maybe Nat])
       (sh :: [Maybe Nat]) i.
proxy sh'
-> StaticShX sh -> ShX (sh ++ sh') i -> (ShX sh i, ShX sh' i)
shxSplitApp (forall (t :: [Maybe Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @sh') StaticShX sh
ssh (StaticShX (sh ++ sh')
-> XArray (sh ++ sh') a -> ShX (sh ++ sh') Int
forall (sh :: [Maybe Nat]) a.
StaticShX sh -> XArray sh a -> IShX sh
shape (StaticShX sh -> StaticShX sh' -> StaticShX (sh ++ sh')
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
StaticShX sh -> StaticShX sh' -> StaticShX (sh ++ sh')
ssxAppend StaticShX sh
ssh StaticShX sh'
ssh') XArray (sh ++ sh') a
arr)
        shF :: ShX '[Flatten' 1 sh] Int
shF = ShX sh Int -> SMayNat Int SNat (Flatten' 1 sh)
forall (sh :: [Maybe Nat]).
IShX sh -> SMayNat Int SNat (Flatten sh)
shxFlatten ShX sh Int
sh SMayNat Int SNat (Flatten' 1 sh)
-> ShX '[] Int -> ShX '[Flatten' 1 sh] Int
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
SMayNat i SNat n -> ShX sh i -> ShX sh1 i
:$% ShX '[] Int
forall (sh :: [Maybe Nat]) i. (sh ~ '[]) => ShX sh i
ZSX
    in StaticShX sh'
-> StaticShX '[Flatten' 1 sh]
-> XArray (sh' ++ '[Flatten' 1 sh]) a
-> XArray sh' a
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) a.
(Storable a, NumElt a) =>
StaticShX sh
-> StaticShX sh' -> XArray (sh ++ sh') a -> XArray sh a
sumInner StaticShX sh'
ssh' (ShX '[Flatten' 1 sh] Int -> StaticShX '[Flatten' 1 sh]
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX ShX '[Flatten' 1 sh] Int
shF) (XArray (sh' ++ '[Flatten' 1 sh]) a -> XArray sh' a)
-> XArray (sh' ++ '[Flatten' 1 sh]) a -> XArray sh' a
forall a b. (a -> b) -> a -> b
$
       StaticShX '[Flatten' 1 sh]
-> StaticShX sh'
-> XArray ('[Flatten' 1 sh] ++ sh') a
-> XArray (sh' ++ '[Flatten' 1 sh]) a
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
StaticShX sh1
-> StaticShX sh2 -> XArray (sh1 ++ sh2) a -> XArray (sh2 ++ sh1) a
transpose2 (ShX '[Flatten' 1 sh] Int -> StaticShX '[Flatten' 1 sh]
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX ShX '[Flatten' 1 sh] Int
shF) StaticShX sh'
ssh' (XArray ('[Flatten' 1 sh] ++ sh') a
 -> XArray (sh' ++ '[Flatten' 1 sh]) a)
-> XArray ('[Flatten' 1 sh] ++ sh') a
-> XArray (sh' ++ '[Flatten' 1 sh]) a
forall a b. (a -> b) -> a -> b
$
       StaticShX sh
-> StaticShX sh'
-> ShX '[Flatten' 1 sh] Int
-> XArray (sh ++ sh') a
-> XArray ('[Flatten' 1 sh] ++ sh') a
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat])
       (sh' :: [Maybe Nat]) a.
Storable a =>
StaticShX sh1
-> StaticShX sh'
-> IShX sh2
-> XArray (sh1 ++ sh') a
-> XArray (sh2 ++ sh') a
reshapePartial StaticShX sh
ssh StaticShX sh'
ssh' ShX '[Flatten' 1 sh] Int
shF (XArray (sh ++ sh') a -> XArray ('[Flatten' 1 sh] ++ sh') a)
-> XArray (sh ++ sh') a -> XArray ('[Flatten' 1 sh] ++ sh') a
forall a b. (a -> b) -> a -> b
$
         XArray (sh ++ sh') a
arr

fromListOuter :: forall n sh a. Storable a
              => StaticShX (n : sh) -> [XArray sh a] -> XArray (n : sh) a
fromListOuter :: forall (n :: Maybe Nat) (sh :: [Maybe Nat]) a.
Storable a =>
StaticShX (n : sh) -> [XArray sh a] -> XArray (n : sh) a
fromListOuter StaticShX (n : sh)
ssh [XArray sh a]
l
  | Dict KnownNat (Rank (n : sh))
Dict <- StaticShX (n : sh) -> Dict KnownNat (Rank (n : sh))
forall (sh :: [Maybe Nat]). StaticShX sh -> Dict KnownNat (Rank sh)
lemKnownNatRankSSX StaticShX (n : sh)
ssh
  = case StaticShX (n : sh)
ssh of
      SKnown SNat n1
m :!% StaticShX sh
_ | SNat n1 -> Int
forall (n :: Nat). SNat n -> Int
fromSNat' SNat n1
m Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [XArray sh a] -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [XArray sh a]
l ->
        String -> XArray (n : sh) a
forall a. HasCallStack => String -> a
error (String -> XArray (n : sh) a) -> String -> XArray (n : sh) a
forall a b. (a -> b) -> a -> b
$ String
"Data.Array.Mixed.fromListOuter: length of list (" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show ([XArray sh a] -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [XArray sh a]
l) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")" String -> ShowS
forall a. [a] -> [a] -> [a]
++
                String
"does not match the type (" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show (SNat n1 -> Int
forall (n :: Nat). SNat n -> Int
fromSNat' SNat n1
m) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"
      StaticShX (n : sh)
_ -> Array (Rank (n : sh)) a -> XArray (n : sh) a
forall (sh :: [Maybe Nat]) a. Array (Rank sh) a -> XArray sh a
XArray (Array 1 (Array (Rank sh) a) -> Array (1 + Rank sh) a
forall a (n :: Nat).
(Unbox a, KnownNat (1 + n)) =>
Array 1 (Array n a) -> Array (1 + n) a
S.ravel ([Int] -> [Array (Rank sh) a] -> Array 1 (Array (Rank sh) a)
forall (n :: Nat) a. KnownNat n => [Int] -> [a] -> Array n a
ORB.fromList [[XArray sh a] -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [XArray sh a]
l] (forall a b. Coercible a b => a -> b
forall a b. Coercible a b => a -> b
coerce @[XArray sh a] @[S.Array (Rank sh) a] [XArray sh a]
l)))

toListOuter :: Storable a => XArray (n : sh) a -> [XArray sh a]
toListOuter :: forall a (n :: Maybe Nat) (sh :: [Maybe Nat]).
Storable a =>
XArray (n : sh) a -> [XArray sh a]
toListOuter (XArray Array (Rank (n : sh)) a
arr) =
  case Array (Rank sh + 1) a -> [Int]
forall (n :: Nat) a. Array n a -> [Int]
S.shapeL Array (Rank sh + 1) a
Array (Rank (n : sh)) a
arr of
    Int
0 : [Int]
_ -> []
    [Int]
_ -> [Array (Rank sh) a] -> [XArray sh a]
forall a b. Coercible a b => a -> b
coerce (Array 1 (Array (Rank sh) a) -> [Array (Rank sh) a]
forall (n :: Nat) a. Array n a -> [a]
ORB.toList (Array (1 + Rank sh) a -> Array 1 (Array (Rank sh) a)
forall a (n :: Nat).
Unbox a =>
Array (1 + n) a -> Array 1 (Array n a)
S.unravel Array (1 + Rank sh) a
Array (Rank (n : sh)) a
arr))

fromList1 :: Storable a => StaticShX '[n] -> [a] -> XArray '[n] a
fromList1 :: forall a (n :: Maybe Nat).
Storable a =>
StaticShX '[n] -> [a] -> XArray '[n] a
fromList1 StaticShX '[n]
ssh [a]
l =
  let n :: Int
n = [a] -> Int
forall a. [a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [a]
l
  in case StaticShX '[n]
ssh of
       SKnown SNat n1
m :!% StaticShX sh
_ | SNat n1 -> Int
forall (n :: Nat). SNat n -> Int
fromSNat' SNat n1
m Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
n ->
         String -> XArray '[n] a
forall a. HasCallStack => String -> a
error (String -> XArray '[n] a) -> String -> XArray '[n] a
forall a b. (a -> b) -> a -> b
$ String
"Data.Array.Mixed.fromList1: length of list (" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
n String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")" String -> ShowS
forall a. [a] -> [a] -> [a]
++
                 String
"does not match the type (" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show (SNat n1 -> Int
forall (n :: Nat). SNat n -> Int
fromSNat' SNat n1
m) String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
")"
       StaticShX '[n]
_ -> Array (Rank '[n]) a -> XArray '[n] a
forall (sh :: [Maybe Nat]) a. Array (Rank sh) a -> XArray sh a
XArray ([Int] -> Vector a -> Array 1 a
forall a (n :: Nat).
(Unbox a, KnownNat n) =>
[Int] -> Vector a -> Array n a
S.fromVector [Int
n] (Int -> [a] -> Vector a
forall a. Storable a => Int -> [a] -> Vector a
VS.fromListN Int
n [a]
l))

toList1 :: Storable a => XArray '[n] a -> [a]
toList1 :: forall a (n :: Maybe Nat). Storable a => XArray '[n] a -> [a]
toList1 (XArray Array (Rank '[n]) a
arr) = Array 1 a -> [a]
forall a (n :: Nat). Unbox a => Array n a -> [a]
S.toList Array 1 a
Array (Rank '[n]) a
arr

-- | Throws if the given shape is not, in fact, empty.
empty :: forall sh a. Storable a => IShX sh -> XArray sh a
empty :: forall (sh :: [Maybe Nat]) a. Storable a => IShX sh -> XArray sh a
empty IShX sh
sh
  | Dict KnownNat (Rank sh)
Dict <- IShX sh -> Dict KnownNat (Rank sh)
forall (sh :: [Maybe Nat]) i. ShX sh i -> Dict KnownNat (Rank sh)
lemKnownNatRank IShX sh
sh
  , IShX sh -> Int
forall (sh :: [Maybe Nat]). IShX sh -> Int
shxSize IShX sh
sh Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
  = Array (Rank sh) a -> XArray sh a
forall (sh :: [Maybe Nat]) a. Array (Rank sh) a -> XArray sh a
XArray ([Int] -> Vector a -> Array (Rank sh) a
forall a (n :: Nat).
(Unbox a, KnownNat n) =>
[Int] -> Vector a -> Array n a
S.fromVector (IShX sh -> [Int]
forall (sh :: [Maybe Nat]). IShX sh -> [Int]
shxToList IShX sh
sh) Vector a
forall a. Storable a => Vector a
VS.empty)
  | Bool
otherwise
  = String -> XArray sh a
forall a. HasCallStack => String -> a
error (String -> XArray sh a) -> String -> XArray sh a
forall a b. (a -> b) -> a -> b
$ String
"Data.Array.Mixed.empty: shape was not empty: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ IShX sh -> String
forall a. Show a => a -> String
show IShX sh
sh

slice :: SNat i -> SNat n -> XArray (Just (i + n + k) : sh) a -> XArray (Just n : sh) a
slice :: forall (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Maybe Nat]) a.
SNat i
-> SNat n
-> XArray ('Just ((i + n) + k) : sh) a
-> XArray ('Just n : sh) a
slice SNat i
i SNat n
n (XArray Array (Rank ('Just ((i + n) + k) : sh)) a
arr) = Array (Rank ('Just n : sh)) a -> XArray ('Just n : sh) a
forall (sh :: [Maybe Nat]) a. Array (Rank sh) a -> XArray sh a
XArray ([(Int, Int)] -> Array (Rank sh + 1) a -> Array (Rank sh + 1) a
forall (n :: Nat) a. [(Int, Int)] -> Array n a -> Array n a
S.slice [(SNat i -> Int
forall (n :: Nat). SNat n -> Int
fromSNat' SNat i
i, SNat n -> Int
forall (n :: Nat). SNat n -> Int
fromSNat' SNat n
n)] Array (Rank sh + 1) a
Array (Rank ('Just ((i + n) + k) : sh)) a
arr)

sliceU :: Int -> Int -> XArray (Nothing : sh) a -> XArray (Nothing : sh) a
sliceU :: forall (sh :: [Maybe Nat]) a.
Int -> Int -> XArray ('Nothing : sh) a -> XArray ('Nothing : sh) a
sliceU Int
i Int
n (XArray Array (Rank ('Nothing : sh)) a
arr) = Array (Rank ('Nothing : sh)) a -> XArray ('Nothing : sh) a
forall (sh :: [Maybe Nat]) a. Array (Rank sh) a -> XArray sh a
XArray ([(Int, Int)] -> Array (Rank sh + 1) a -> Array (Rank sh + 1) a
forall (n :: Nat) a. [(Int, Int)] -> Array n a -> Array n a
S.slice [(Int
i, Int
n)] Array (Rank sh + 1) a
Array (Rank ('Nothing : sh)) a
arr)

rev1 :: XArray (n : sh) a -> XArray (n : sh) a
rev1 :: forall (n :: Maybe Nat) (sh :: [Maybe Nat]) a.
XArray (n : sh) a -> XArray (n : sh) a
rev1 (XArray Array (Rank (n : sh)) a
arr) = Array (Rank (n : sh)) a -> XArray (n : sh) a
forall (sh :: [Maybe Nat]) a. Array (Rank sh) a -> XArray sh a
XArray ([Int] -> Array (Rank sh + 1) a -> Array (Rank sh + 1) a
forall (n :: Nat) a. [Int] -> Array n a -> Array n a
S.rev [Int
0] Array (Rank sh + 1) a
Array (Rank (n : sh)) a
arr)

-- | Throws if the given array and the target shape do not have the same number of elements.
reshape :: forall sh1 sh2 a. Storable a => StaticShX sh1 -> IShX sh2 -> XArray sh1 a -> XArray sh2 a
reshape :: forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
Storable a =>
StaticShX sh1 -> IShX sh2 -> XArray sh1 a -> XArray sh2 a
reshape StaticShX sh1
ssh1 IShX sh2
sh2 (XArray Array (Rank sh1) a
arr)
  | Dict KnownNat (Rank sh1)
Dict <- StaticShX sh1 -> Dict KnownNat (Rank sh1)
forall (sh :: [Maybe Nat]). StaticShX sh -> Dict KnownNat (Rank sh)
lemKnownNatRankSSX StaticShX sh1
ssh1
  , Dict KnownNat (Rank sh2)
Dict <- IShX sh2 -> Dict KnownNat (Rank sh2)
forall (sh :: [Maybe Nat]) i. ShX sh i -> Dict KnownNat (Rank sh)
lemKnownNatRank IShX sh2
sh2
  = Array (Rank sh2) a -> XArray sh2 a
forall (sh :: [Maybe Nat]) a. Array (Rank sh) a -> XArray sh a
XArray ([Int] -> Array (Rank sh1) a -> Array (Rank sh2) a
forall a (n :: Nat) (n' :: Nat).
(Unbox a, KnownNat n, KnownNat n') =>
[Int] -> Array n a -> Array n' a
S.reshape (IShX sh2 -> [Int]
forall (sh :: [Maybe Nat]). IShX sh -> [Int]
shxToList IShX sh2
sh2) Array (Rank sh1) a
arr)

-- | Throws if the given array and the target shape do not have the same number of elements.
reshapePartial :: forall sh1 sh2 sh' a. Storable a => StaticShX sh1 -> StaticShX sh' -> IShX sh2 -> XArray (sh1 ++ sh') a -> XArray (sh2 ++ sh') a
reshapePartial :: forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat])
       (sh' :: [Maybe Nat]) a.
Storable a =>
StaticShX sh1
-> StaticShX sh'
-> IShX sh2
-> XArray (sh1 ++ sh') a
-> XArray (sh2 ++ sh') a
reshapePartial StaticShX sh1
ssh1 StaticShX sh'
ssh' IShX sh2
sh2 (XArray Array (Rank (sh1 ++ sh')) a
arr)
  | Dict KnownNat (Rank (sh1 ++ sh'))
Dict <- StaticShX (sh1 ++ sh') -> Dict KnownNat (Rank (sh1 ++ sh'))
forall (sh :: [Maybe Nat]). StaticShX sh -> Dict KnownNat (Rank sh)
lemKnownNatRankSSX (StaticShX sh1 -> StaticShX sh' -> StaticShX (sh1 ++ sh')
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
StaticShX sh -> StaticShX sh' -> StaticShX (sh ++ sh')
ssxAppend StaticShX sh1
ssh1 StaticShX sh'
ssh')
  , Dict KnownNat (Rank (sh2 ++ sh'))
Dict <- StaticShX (sh2 ++ sh') -> Dict KnownNat (Rank (sh2 ++ sh'))
forall (sh :: [Maybe Nat]). StaticShX sh -> Dict KnownNat (Rank sh)
lemKnownNatRankSSX (StaticShX sh2 -> StaticShX sh' -> StaticShX (sh2 ++ sh')
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
StaticShX sh -> StaticShX sh' -> StaticShX (sh ++ sh')
ssxAppend (IShX sh2 -> StaticShX sh2
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX IShX sh2
sh2) StaticShX sh'
ssh')
  = Array (Rank (sh2 ++ sh')) a -> XArray (sh2 ++ sh') a
forall (sh :: [Maybe Nat]) a. Array (Rank sh) a -> XArray sh a
XArray ([Int] -> Array (Rank (sh1 ++ sh')) a -> Array (Rank (sh2 ++ sh')) a
forall a (n :: Nat) (n' :: Nat).
(Unbox a, KnownNat n, KnownNat n') =>
[Int] -> Array n a -> Array n' a
S.reshape (IShX sh2 -> [Int]
forall (sh :: [Maybe Nat]). IShX sh -> [Int]
shxToList IShX sh2
sh2 [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
drop (StaticShX sh1 -> Int
forall (sh :: [Maybe Nat]). StaticShX sh -> Int
ssxLength StaticShX sh1
ssh1) (Array (Rank (sh1 ++ sh')) a -> [Int]
forall (n :: Nat) a. Array n a -> [Int]
S.shapeL Array (Rank (sh1 ++ sh')) a
arr)) Array (Rank (sh1 ++ sh')) a
arr)

-- this was benchmarked to be (slightly) faster than S.iota, S.generate and S.fromVector(VS.enumFromTo).
iota :: (Enum a, Storable a) => SNat n -> XArray '[Just n] a
iota :: forall a (n :: Nat).
(Enum a, Storable a) =>
SNat n -> XArray '[ 'Just n] a
iota SNat n
sn = Array (Rank '[ 'Just n]) a -> XArray '[ 'Just n] a
forall (sh :: [Maybe Nat]) a. Array (Rank sh) a -> XArray sh a
XArray ([Int] -> Vector a -> Array 1 a
forall a (n :: Nat).
(Unbox a, KnownNat n) =>
[Int] -> Vector a -> Array n a
S.fromVector [SNat n -> Int
forall (n :: Nat). SNat n -> Int
fromSNat' SNat n
sn] (Int -> [a] -> Vector a
forall a. Storable a => Int -> [a] -> Vector a
VS.fromListN (SNat n -> Int
forall (n :: Nat). SNat n -> Int
fromSNat' SNat n
sn) [Int -> a
forall a. Enum a => Int -> a
toEnum Int
0 .. Int -> a
forall a. Enum a => Int -> a
toEnum (SNat n -> Int
forall (n :: Nat). SNat n -> Int
fromSNat' SNat n
sn Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)]))