{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ImportQualifiedPost #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE ViewPatterns #-}
module Data.Array.Nested.Shaped (
  Shaped(Shaped),
  squotArray, sremArray, satan2Array,
  sshape,
  module Data.Array.Nested.Shaped,
  liftShaped1, liftShaped2,
) where

import Prelude hiding (mappend, mconcat)

import Data.Array.Internal.RankedG qualified as RG
import Data.Array.Internal.RankedS qualified as RS
import Data.Array.Internal.ShapedG qualified as SG
import Data.Array.Internal.ShapedS qualified as SS
import Data.Bifunctor (first)
import Data.Coerce (coerce)
import Data.List.NonEmpty (NonEmpty)
import Data.Proxy
import Data.Type.Equality
import Data.Vector.Storable qualified as VS
import Foreign.Storable (Storable)
import GHC.TypeLits

import Data.Array.Nested.Convert
import Data.Array.Nested.Lemmas
import Data.Array.Nested.Mixed
import Data.Array.Nested.Mixed.Shape
import Data.Array.Nested.Permutation
import Data.Array.Nested.Shaped.Base
import Data.Array.Nested.Shaped.Shape
import Data.Array.Nested.Types
import Data.Array.Strided.Arith
import Data.Array.XArray (XArray)
import Data.Array.XArray qualified as X


semptyArray :: KnownElt a => ShS sh -> Shaped (0 : sh) a
semptyArray :: forall a (sh :: [Nat]). KnownElt a => ShS sh -> Shaped (0 : sh) a
semptyArray ShS sh
sh = Mixed (MapJust (0 : sh)) a -> Shaped (0 : sh) a
forall (sh :: [Nat]) a. Mixed (MapJust sh) a -> Shaped sh a
Shaped (IShX (MapJust sh) -> Mixed ('Just 0 : MapJust sh) a
forall a (sh :: [Maybe Nat]).
KnownElt a =>
IShX sh -> Mixed ('Just 0 : sh) a
memptyArray (ShS sh -> IShX (MapJust sh)
forall (sh :: [Nat]). ShS sh -> IShX (MapJust sh)
shxFromShS ShS sh
sh))

srank :: Elt a => Shaped sh a -> SNat (Rank sh)
srank :: forall a (sh :: [Nat]). Elt a => Shaped sh a -> SNat (Rank sh)
srank = ShS sh -> SNat (Rank sh)
forall (sh :: [Nat]). ShS sh -> SNat (Rank sh)
shsRank (ShS sh -> SNat (Rank sh))
-> (Shaped sh a -> ShS sh) -> Shaped sh a -> SNat (Rank sh)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shaped sh a -> ShS sh
forall (sh :: [Nat]) a. Elt a => Shaped sh a -> ShS sh
sshape

-- | The total number of elements in the array.
ssize :: Elt a => Shaped sh a -> Int
ssize :: forall a (sh :: [Nat]). Elt a => Shaped sh a -> Int
ssize = ShS sh -> Int
forall (sh :: [Nat]). ShS sh -> Int
shsSize (ShS sh -> Int) -> (Shaped sh a -> ShS sh) -> Shaped sh a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shaped sh a -> ShS sh
forall (sh :: [Nat]) a. Elt a => Shaped sh a -> ShS sh
sshape

sindex :: Elt a => Shaped sh a -> IIxS sh -> a
sindex :: forall a (sh :: [Nat]). Elt a => Shaped sh a -> IIxS sh -> a
sindex (Shaped Mixed (MapJust sh) a
arr) IIxS sh
idx = Mixed (MapJust sh) a -> IIxX (MapJust sh) -> a
forall (sh :: [Maybe Nat]). Mixed sh a -> IIxX sh -> a
forall a (sh :: [Maybe Nat]). Elt a => Mixed sh a -> IIxX sh -> a
mindex Mixed (MapJust sh) a
arr (IIxS sh -> IIxX (MapJust sh)
forall (sh :: [Nat]) i. IxS sh i -> IxX (MapJust sh) i
ixxFromIxS IIxS sh
idx)

shsTakeIx :: Proxy sh' -> ShS (sh ++ sh') -> IIxS sh -> ShS sh
shsTakeIx :: forall (sh' :: [Nat]) (sh :: [Nat]).
Proxy sh' -> ShS (sh ++ sh') -> IIxS sh -> ShS sh
shsTakeIx Proxy sh'
_ ShS (sh ++ sh')
_ IxS sh Int
ZIS = ShS sh
forall (sh :: [Nat]). (sh ~ '[]) => ShS sh
ZSS
shsTakeIx Proxy sh'
p ShS (sh ++ sh')
sh (Int
_ :.$ IxS sh Int
idx) = case ShS (sh ++ sh')
sh of SNat n
n :$$ ShS sh
sh' -> SNat n
n SNat n -> ShS sh -> ShS sh
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, (n : sh) ~ sh1) =>
SNat n -> ShS sh -> ShS sh1
:$$ Proxy sh' -> ShS (sh ++ sh') -> IxS sh Int -> ShS sh
forall (sh' :: [Nat]) (sh :: [Nat]).
Proxy sh' -> ShS (sh ++ sh') -> IIxS sh -> ShS sh
shsTakeIx Proxy sh'
p ShS sh
ShS (sh ++ sh')
sh' IxS sh Int
idx

sindexPartial :: forall sh1 sh2 a. Elt a => Shaped (sh1 ++ sh2) a -> IIxS sh1 -> Shaped sh2 a
sindexPartial :: forall (sh1 :: [Nat]) (sh2 :: [Nat]) a.
Elt a =>
Shaped (sh1 ++ sh2) a -> IIxS sh1 -> Shaped sh2 a
sindexPartial sarr :: Shaped (sh1 ++ sh2) a
sarr@(Shaped Mixed (MapJust (sh1 ++ sh2)) a
arr) IIxS sh1
idx =
  Mixed (MapJust sh2) a -> Shaped sh2 a
forall (sh :: [Nat]) a. Mixed (MapJust sh) a -> Shaped sh a
Shaped (forall a (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
Elt a =>
Mixed (sh ++ sh') a -> IIxX sh -> Mixed sh' a
mindexPartial @a @(MapJust sh1) @(MapJust sh2)
            ((Mixed (MapJust (sh1 ++ sh2)) a
 :~: Mixed (MapJust sh1 ++ MapJust sh2) a)
-> Mixed (MapJust (sh1 ++ sh2)) a
-> Mixed (MapJust sh1 ++ MapJust sh2) a
forall a b. (a :~: b) -> a -> b
castWith ((MapJust (sh1 ++ sh2) :~: (MapJust sh1 ++ MapJust sh2))
-> Mixed (MapJust (sh1 ++ sh2)) a
   :~: Mixed (MapJust sh1 ++ MapJust sh2) a
forall {k1} {k2} {k3} (f :: k1 -> k2 -> k3) (c :: k2) (a :: k1)
       (b :: k1).
(a :~: b) -> f a c :~: f b c
subst2 (ShS sh1
-> Proxy sh2
-> MapJust (sh1 ++ sh2) :~: (MapJust sh1 ++ MapJust sh2)
forall (sh1 :: [Nat]) (sh2 :: [Nat]).
ShS sh1
-> Proxy sh2
-> MapJust (sh1 ++ sh2) :~: (MapJust sh1 ++ MapJust sh2)
lemMapJustApp (Proxy sh2 -> ShS (sh1 ++ sh2) -> IIxS sh1 -> ShS sh1
forall (sh' :: [Nat]) (sh :: [Nat]).
Proxy sh' -> ShS (sh ++ sh') -> IIxS sh -> ShS sh
shsTakeIx (forall (t :: [Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @sh2) (Shaped (sh1 ++ sh2) a -> ShS (sh1 ++ sh2)
forall (sh :: [Nat]) a. Elt a => Shaped sh a -> ShS sh
sshape Shaped (sh1 ++ sh2) a
sarr) IIxS sh1
idx) (forall (t :: [Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @sh2))) Mixed (MapJust (sh1 ++ sh2)) a
arr)
            (IIxS sh1 -> IxX (MapJust sh1) Int
forall (sh :: [Nat]) i. IxS sh i -> IxX (MapJust sh) i
ixxFromIxS IIxS sh1
idx))

-- | __WARNING__: All values returned from the function must have equal shape.
-- See the documentation of 'mgenerate' for more details.
sgenerate :: forall sh a. KnownElt a => ShS sh -> (IIxS sh -> a) -> Shaped sh a
sgenerate :: forall (sh :: [Nat]) a.
KnownElt a =>
ShS sh -> (IIxS sh -> a) -> Shaped sh a
sgenerate ShS sh
sh IIxS sh -> a
f = Mixed (MapJust sh) a -> Shaped sh a
forall (sh :: [Nat]) a. Mixed (MapJust sh) a -> Shaped sh a
Shaped (IShX (MapJust sh)
-> (IIxX (MapJust sh) -> a) -> Mixed (MapJust sh) a
forall (sh :: [Maybe Nat]) a.
KnownElt a =>
IShX sh -> (IIxX sh -> a) -> Mixed sh a
mgenerate (ShS sh -> IShX (MapJust sh)
forall (sh :: [Nat]). ShS sh -> IShX (MapJust sh)
shxFromShS ShS sh
sh) (IIxS sh -> a
f (IIxS sh -> a)
-> (IIxX (MapJust sh) -> IIxS sh) -> IIxX (MapJust sh) -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShS sh -> IIxX (MapJust sh) -> IIxS sh
forall (sh :: [Nat]) i. ShS sh -> IxX (MapJust sh) i -> IxS sh i
ixsFromIxX ShS sh
sh))

-- | See the documentation of 'mlift'.
slift :: forall sh1 sh2 a. Elt a
      => ShS sh2
      -> (forall sh' b. Storable b => StaticShX sh' -> XArray (MapJust sh1 ++ sh') b -> XArray (MapJust sh2 ++ sh') b)
      -> Shaped sh1 a -> Shaped sh2 a
slift :: forall (sh1 :: [Nat]) (sh2 :: [Nat]) a.
Elt a =>
ShS sh2
-> (forall (sh' :: [Maybe Nat]) b.
    Storable b =>
    StaticShX sh'
    -> XArray (MapJust sh1 ++ sh') b -> XArray (MapJust sh2 ++ sh') b)
-> Shaped sh1 a
-> Shaped sh2 a
slift ShS sh2
sh2 forall (sh' :: [Maybe Nat]) b.
Storable b =>
StaticShX sh'
-> XArray (MapJust sh1 ++ sh') b -> XArray (MapJust sh2 ++ sh') b
f (Shaped Mixed (MapJust sh1) a
arr) = Mixed (MapJust sh2) a -> Shaped sh2 a
forall (sh :: [Nat]) a. Mixed (MapJust sh) a -> Shaped sh a
Shaped (StaticShX (MapJust sh2)
-> (forall (sh' :: [Maybe Nat]) b.
    Storable b =>
    StaticShX sh'
    -> XArray (MapJust sh1 ++ sh') b -> XArray (MapJust sh2 ++ sh') b)
-> Mixed (MapJust sh1) a
-> Mixed (MapJust sh2) a
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]).
StaticShX sh2
-> (forall (sh' :: [Maybe Nat]) b.
    Storable b =>
    StaticShX sh' -> XArray (sh1 ++ sh') b -> XArray (sh2 ++ sh') b)
-> Mixed sh1 a
-> Mixed sh2 a
forall a (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]).
Elt a =>
StaticShX sh2
-> (forall (sh' :: [Maybe Nat]) b.
    Storable b =>
    StaticShX sh' -> XArray (sh1 ++ sh') b -> XArray (sh2 ++ sh') b)
-> Mixed sh1 a
-> Mixed sh2 a
mlift (ShX (MapJust sh2) Int -> StaticShX (MapJust sh2)
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX (ShS sh2 -> ShX (MapJust sh2) Int
forall (sh :: [Nat]). ShS sh -> IShX (MapJust sh)
shxFromShS ShS sh2
sh2)) StaticShX sh'
-> XArray (MapJust sh1 ++ sh') b -> XArray (MapJust sh2 ++ sh') b
forall (sh' :: [Maybe Nat]) b.
Storable b =>
StaticShX sh'
-> XArray (MapJust sh1 ++ sh') b -> XArray (MapJust sh2 ++ sh') b
f Mixed (MapJust sh1) a
arr)

-- | See the documentation of 'mlift'.
slift2 :: forall sh1 sh2 sh3 a. Elt a
       => ShS sh3
       -> (forall sh' b. Storable b => StaticShX sh' -> XArray (MapJust sh1 ++ sh') b -> XArray (MapJust sh2 ++ sh') b -> XArray (MapJust sh3 ++ sh') b)
       -> Shaped sh1 a -> Shaped sh2 a -> Shaped sh3 a
slift2 :: forall (sh1 :: [Nat]) (sh2 :: [Nat]) (sh3 :: [Nat]) a.
Elt a =>
ShS sh3
-> (forall (sh' :: [Maybe Nat]) b.
    Storable b =>
    StaticShX sh'
    -> XArray (MapJust sh1 ++ sh') b
    -> XArray (MapJust sh2 ++ sh') b
    -> XArray (MapJust sh3 ++ sh') b)
-> Shaped sh1 a
-> Shaped sh2 a
-> Shaped sh3 a
slift2 ShS sh3
sh3 forall (sh' :: [Maybe Nat]) b.
Storable b =>
StaticShX sh'
-> XArray (MapJust sh1 ++ sh') b
-> XArray (MapJust sh2 ++ sh') b
-> XArray (MapJust sh3 ++ sh') b
f (Shaped Mixed (MapJust sh1) a
arr1) (Shaped Mixed (MapJust sh2) a
arr2) = Mixed (MapJust sh3) a -> Shaped sh3 a
forall (sh :: [Nat]) a. Mixed (MapJust sh) a -> Shaped sh a
Shaped (StaticShX (MapJust sh3)
-> (forall (sh' :: [Maybe Nat]) b.
    Storable b =>
    StaticShX sh'
    -> XArray (MapJust sh1 ++ sh') b
    -> XArray (MapJust sh2 ++ sh') b
    -> XArray (MapJust sh3 ++ sh') b)
-> Mixed (MapJust sh1) a
-> Mixed (MapJust sh2) a
-> Mixed (MapJust sh3) a
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat])
       (sh3 :: [Maybe Nat]).
StaticShX sh3
-> (forall (sh' :: [Maybe Nat]) b.
    Storable b =>
    StaticShX sh'
    -> XArray (sh1 ++ sh') b
    -> XArray (sh2 ++ sh') b
    -> XArray (sh3 ++ sh') b)
-> Mixed sh1 a
-> Mixed sh2 a
-> Mixed sh3 a
forall a (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat])
       (sh3 :: [Maybe Nat]).
Elt a =>
StaticShX sh3
-> (forall (sh' :: [Maybe Nat]) b.
    Storable b =>
    StaticShX sh'
    -> XArray (sh1 ++ sh') b
    -> XArray (sh2 ++ sh') b
    -> XArray (sh3 ++ sh') b)
-> Mixed sh1 a
-> Mixed sh2 a
-> Mixed sh3 a
mlift2 (ShX (MapJust sh3) Int -> StaticShX (MapJust sh3)
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX (ShS sh3 -> ShX (MapJust sh3) Int
forall (sh :: [Nat]). ShS sh -> IShX (MapJust sh)
shxFromShS ShS sh3
sh3)) StaticShX sh'
-> XArray (MapJust sh1 ++ sh') b
-> XArray (MapJust sh2 ++ sh') b
-> XArray (MapJust sh3 ++ sh') b
forall (sh' :: [Maybe Nat]) b.
Storable b =>
StaticShX sh'
-> XArray (MapJust sh1 ++ sh') b
-> XArray (MapJust sh2 ++ sh') b
-> XArray (MapJust sh3 ++ sh') b
f Mixed (MapJust sh1) a
arr1 Mixed (MapJust sh2) a
arr2)

ssumOuter1P :: forall sh n a. (Storable a, NumElt a)
            => Shaped (n : sh) (Primitive a) -> Shaped sh (Primitive a)
ssumOuter1P :: forall (sh :: [Nat]) (n :: Nat) a.
(Storable a, NumElt a) =>
Shaped (n : sh) (Primitive a) -> Shaped sh (Primitive a)
ssumOuter1P (Shaped Mixed (MapJust (n : sh)) (Primitive a)
arr) = Mixed (MapJust sh) (Primitive a) -> Shaped sh (Primitive a)
forall (sh :: [Nat]) a. Mixed (MapJust sh) a -> Shaped sh a
Shaped (Mixed ('Just n : MapJust sh) (Primitive a)
-> Mixed (MapJust sh) (Primitive a)
forall (sh :: [Maybe Nat]) (n :: Maybe Nat) a.
(Storable a, NumElt a) =>
Mixed (n : sh) (Primitive a) -> Mixed sh (Primitive a)
msumOuter1P Mixed ('Just n : MapJust sh) (Primitive a)
Mixed (MapJust (n : sh)) (Primitive a)
arr)

ssumOuter1 :: forall sh n a. (NumElt a, PrimElt a)
           => Shaped (n : sh) a -> Shaped sh a
ssumOuter1 :: forall (sh :: [Nat]) (n :: Nat) a.
(NumElt a, PrimElt a) =>
Shaped (n : sh) a -> Shaped sh a
ssumOuter1 = Shaped sh (Primitive a) -> Shaped sh a
forall a (sh :: [Nat]).
PrimElt a =>
Shaped sh (Primitive a) -> Shaped sh a
sfromPrimitive (Shaped sh (Primitive a) -> Shaped sh a)
-> (Shaped (n : sh) a -> Shaped sh (Primitive a))
-> Shaped (n : sh) a
-> Shaped sh a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shaped (n : sh) (Primitive a) -> Shaped sh (Primitive a)
forall (sh :: [Nat]) (n :: Nat) a.
(Storable a, NumElt a) =>
Shaped (n : sh) (Primitive a) -> Shaped sh (Primitive a)
ssumOuter1P (Shaped (n : sh) (Primitive a) -> Shaped sh (Primitive a))
-> (Shaped (n : sh) a -> Shaped (n : sh) (Primitive a))
-> Shaped (n : sh) a
-> Shaped sh (Primitive a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shaped (n : sh) a -> Shaped (n : sh) (Primitive a)
forall a (sh :: [Nat]).
PrimElt a =>
Shaped sh a -> Shaped sh (Primitive a)
stoPrimitive

ssumAllPrim :: (PrimElt a, NumElt a) => Shaped n a -> a
ssumAllPrim :: forall a (n :: [Nat]). (PrimElt a, NumElt a) => Shaped n a -> a
ssumAllPrim (Shaped Mixed (MapJust n) a
arr) = Mixed (MapJust n) a -> a
forall a (sh :: [Maybe Nat]).
(PrimElt a, NumElt a) =>
Mixed sh a -> a
msumAllPrim Mixed (MapJust n) a
arr

stranspose :: forall is sh a. (IsPermutation is, Rank is <= Rank sh, Elt a)
           => Perm is -> Shaped sh a -> Shaped (PermutePrefix is sh) a
stranspose :: forall (is :: [Nat]) (sh :: [Nat]) a.
(IsPermutation is, Rank is <= Rank sh, Elt a) =>
Perm is -> Shaped sh a -> Shaped (PermutePrefix is sh) a
stranspose Perm is
perm sarr :: Shaped sh a
sarr@(Shaped Mixed (MapJust sh) a
arr)
  | Rank (MapJust sh) :~: Rank sh
Refl <- ShS sh -> Rank (MapJust sh) :~: Rank sh
forall (sh :: [Nat]). ShS sh -> Rank (MapJust sh) :~: Rank sh
lemRankMapJust (Shaped sh a -> ShS sh
forall (sh :: [Nat]) a. Elt a => Shaped sh a -> ShS sh
sshape Shaped sh a
sarr)
  , TakeLen is (MapJust sh) :~: MapJust (TakeLen is sh)
Refl <- Perm is
-> ShS sh -> TakeLen is (MapJust sh) :~: MapJust (TakeLen is sh)
forall (is :: [Nat]) (sh :: [Nat]).
Perm is
-> ShS sh -> TakeLen is (MapJust sh) :~: MapJust (TakeLen is sh)
lemTakeLenMapJust Perm is
perm (Shaped sh a -> ShS sh
forall (sh :: [Nat]) a. Elt a => Shaped sh a -> ShS sh
sshape Shaped sh a
sarr)
  , DropLen is (MapJust sh) :~: MapJust (DropLen is sh)
Refl <- Perm is
-> ShS sh -> DropLen is (MapJust sh) :~: MapJust (DropLen is sh)
forall (is :: [Nat]) (sh :: [Nat]).
Perm is
-> ShS sh -> DropLen is (MapJust sh) :~: MapJust (DropLen is sh)
lemDropLenMapJust Perm is
perm (Shaped sh a -> ShS sh
forall (sh :: [Nat]) a. Elt a => Shaped sh a -> ShS sh
sshape Shaped sh a
sarr)
  , Permute is (MapJust (TakeLen is sh))
:~: MapJust (Permute is (TakeLen is sh))
Refl <- Perm is
-> ShS (TakeLen is sh)
-> Permute is (MapJust (TakeLen is sh))
   :~: MapJust (Permute is (TakeLen is sh))
forall (is :: [Nat]) (sh :: [Nat]).
Perm is
-> ShS sh -> Permute is (MapJust sh) :~: MapJust (Permute is sh)
lemPermuteMapJust Perm is
perm (Perm is -> ShS sh -> ShS (TakeLen is sh)
forall (is :: [Nat]) (sh :: [Nat]).
Perm is -> ShS sh -> ShS (TakeLen is sh)
shsTakeLen Perm is
perm (Shaped sh a -> ShS sh
forall (sh :: [Nat]) a. Elt a => Shaped sh a -> ShS sh
sshape Shaped sh a
sarr))
  , MapJust (Permute is (TakeLen is sh) ++ DropLen is sh)
:~: (MapJust (Permute is (TakeLen is sh))
     ++ MapJust (DropLen is sh))
Refl <- ShS (Permute is (TakeLen is sh))
-> Proxy (DropLen is sh)
-> MapJust (Permute is (TakeLen is sh) ++ DropLen is sh)
   :~: (MapJust (Permute is (TakeLen is sh))
        ++ MapJust (DropLen is sh))
forall (sh1 :: [Nat]) (sh2 :: [Nat]).
ShS sh1
-> Proxy sh2
-> MapJust (sh1 ++ sh2) :~: (MapJust sh1 ++ MapJust sh2)
lemMapJustApp (Perm is -> ShS (TakeLen is sh) -> ShS (Permute is (TakeLen is sh))
forall (is :: [Nat]) (sh :: [Nat]).
Perm is -> ShS sh -> ShS (Permute is sh)
shsPermute Perm is
perm (Perm is -> ShS sh -> ShS (TakeLen is sh)
forall (is :: [Nat]) (sh :: [Nat]).
Perm is -> ShS sh -> ShS (TakeLen is sh)
shsTakeLen Perm is
perm (Shaped sh a -> ShS sh
forall (sh :: [Nat]) a. Elt a => Shaped sh a -> ShS sh
sshape Shaped sh a
sarr))) (forall (t :: [Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @(DropLen is sh))
  = Mixed (MapJust (Permute is (TakeLen is sh) ++ DropLen is sh)) a
-> Shaped (Permute is (TakeLen is sh) ++ DropLen is sh) a
forall (sh :: [Nat]) a. Mixed (MapJust sh) a -> Shaped sh a
Shaped (Perm is
-> Mixed (MapJust sh) a -> Mixed (PermutePrefix is (MapJust sh)) a
forall (is :: [Nat]) (sh :: [Maybe Nat]).
(IsPermutation is, Rank is <= Rank sh) =>
Perm is -> Mixed sh a -> Mixed (PermutePrefix is sh) a
forall a (is :: [Nat]) (sh :: [Maybe Nat]).
(Elt a, IsPermutation is, Rank is <= Rank sh) =>
Perm is -> Mixed sh a -> Mixed (PermutePrefix is sh) a
mtranspose Perm is
perm Mixed (MapJust sh) a
arr)

sappend :: Elt a => Shaped (n : sh) a -> Shaped (m : sh) a -> Shaped (n + m : sh) a
sappend :: forall a (n :: Nat) (sh :: [Nat]) (m :: Nat).
Elt a =>
Shaped (n : sh) a -> Shaped (m : sh) a -> Shaped ((n + m) : sh) a
sappend = (Mixed ('Just n : MapJust sh) a
 -> Mixed ('Just m : MapJust sh) a
 -> Mixed (AddMaybe ('Just n) ('Just m) : MapJust sh) a)
-> Shaped (n : sh) a
-> Shaped (m : sh) a
-> Shaped ((n + m) : sh) a
forall a b. Coercible a b => a -> b
coerce Mixed ('Just n : MapJust sh) a
-> Mixed ('Just m : MapJust sh) a
-> Mixed (AddMaybe ('Just n) ('Just m) : MapJust sh) a
forall (n :: Maybe Nat) (m :: Maybe Nat) (sh :: [Maybe Nat]) a.
Elt a =>
Mixed (n : sh) a -> Mixed (m : sh) a -> Mixed (AddMaybe n m : sh) a
mappend

sscalar :: Elt a => a -> Shaped '[] a
sscalar :: forall a. Elt a => a -> Shaped '[] a
sscalar a
x = Mixed (MapJust '[]) a -> Shaped '[] a
forall (sh :: [Nat]) a. Mixed (MapJust sh) a -> Shaped sh a
Shaped (a -> Mixed '[] a
forall a. Elt a => a -> Mixed '[] a
mscalar a
x)

sfromVectorP :: Storable a => ShS sh -> VS.Vector a -> Shaped sh (Primitive a)
sfromVectorP :: forall a (sh :: [Nat]).
Storable a =>
ShS sh -> Vector a -> Shaped sh (Primitive a)
sfromVectorP ShS sh
sh Vector a
v = Mixed (MapJust sh) (Primitive a) -> Shaped sh (Primitive a)
forall (sh :: [Nat]) a. Mixed (MapJust sh) a -> Shaped sh a
Shaped (IShX (MapJust sh) -> Vector a -> Mixed (MapJust sh) (Primitive a)
forall (sh :: [Maybe Nat]) a.
Storable a =>
IShX sh -> Vector a -> Mixed sh (Primitive a)
mfromVectorP (ShS sh -> IShX (MapJust sh)
forall (sh :: [Nat]). ShS sh -> IShX (MapJust sh)
shxFromShS ShS sh
sh) Vector a
v)

sfromVector :: PrimElt a => ShS sh -> VS.Vector a -> Shaped sh a
sfromVector :: forall a (sh :: [Nat]).
PrimElt a =>
ShS sh -> Vector a -> Shaped sh a
sfromVector ShS sh
sh Vector a
v = Shaped sh (Primitive a) -> Shaped sh a
forall a (sh :: [Nat]).
PrimElt a =>
Shaped sh (Primitive a) -> Shaped sh a
sfromPrimitive (ShS sh -> Vector a -> Shaped sh (Primitive a)
forall a (sh :: [Nat]).
Storable a =>
ShS sh -> Vector a -> Shaped sh (Primitive a)
sfromVectorP ShS sh
sh Vector a
v)

stoVectorP :: Storable a => Shaped sh (Primitive a) -> VS.Vector a
stoVectorP :: forall a (sh :: [Nat]).
Storable a =>
Shaped sh (Primitive a) -> Vector a
stoVectorP = (Mixed (MapJust sh) (Primitive a) -> Vector a)
-> Shaped sh (Primitive a) -> Vector a
forall a b. Coercible a b => a -> b
coerce Mixed (MapJust sh) (Primitive a) -> Vector a
forall a (sh :: [Maybe Nat]).
Storable a =>
Mixed sh (Primitive a) -> Vector a
mtoVectorP

stoVector :: PrimElt a => Shaped sh a -> VS.Vector a
stoVector :: forall a (sh :: [Nat]). PrimElt a => Shaped sh a -> Vector a
stoVector = (Mixed (MapJust sh) a -> Vector a) -> Shaped sh a -> Vector a
forall a b. Coercible a b => a -> b
coerce Mixed (MapJust sh) a -> Vector a
forall a (sh :: [Maybe Nat]). PrimElt a => Mixed sh a -> Vector a
mtoVector

sfromList1 :: Elt a => SNat n -> NonEmpty a -> Shaped '[n] a
sfromList1 :: forall a (n :: Nat). Elt a => SNat n -> NonEmpty a -> Shaped '[n] a
sfromList1 SNat n
sn = Mixed '[ 'Just n] a -> Shaped '[n] a
Mixed (MapJust '[n]) a -> Shaped '[n] a
forall (sh :: [Nat]) a. Mixed (MapJust sh) a -> Shaped sh a
Shaped (Mixed '[ 'Just n] a -> Shaped '[n] a)
-> (NonEmpty a -> Mixed '[ 'Just n] a)
-> NonEmpty a
-> Shaped '[n] a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StaticShX '[ 'Just n]
-> Mixed '[ 'Nothing] a -> Mixed '[ 'Just n] a
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
(Rank sh1 ~ Rank sh2, Elt a) =>
StaticShX sh2 -> Mixed sh1 a -> Mixed sh2 a
mcast (SNat n -> SMayNat () SNat ('Just n)
forall {k} (f :: k -> *) (n1 :: k) i.
f n1 -> SMayNat i f ('Just n1)
SKnown SNat n
sn SMayNat () SNat ('Just n) -> StaticShX '[] -> StaticShX '[ 'Just n]
forall {sh1 :: [Maybe Nat]} (n :: Maybe Nat) (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
SMayNat () SNat n -> StaticShX sh -> StaticShX sh1
:!% StaticShX '[]
forall (sh :: [Maybe Nat]). (sh ~ '[]) => StaticShX sh
ZKX) (Mixed '[ 'Nothing] a -> Mixed '[ 'Just n] a)
-> (NonEmpty a -> Mixed '[ 'Nothing] a)
-> NonEmpty a
-> Mixed '[ 'Just n] a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. NonEmpty a -> Mixed '[ 'Nothing] a
forall a. Elt a => NonEmpty a -> Mixed '[ 'Nothing] a
mfromList1

sfromListOuter :: Elt a => SNat n -> NonEmpty (Shaped sh a) -> Shaped (n : sh) a
sfromListOuter :: forall a (n :: Nat) (sh :: [Nat]).
Elt a =>
SNat n -> NonEmpty (Shaped sh a) -> Shaped (n : sh) a
sfromListOuter SNat n
sn NonEmpty (Shaped sh a)
l = Mixed (MapJust (n : sh)) a -> Shaped (n : sh) a
forall (sh :: [Nat]) a. Mixed (MapJust sh) a -> Shaped sh a
Shaped (StaticShX '[ 'Nothing]
-> StaticShX '[ 'Just n]
-> Proxy (MapJust sh)
-> Mixed ('[ 'Nothing] ++ MapJust sh) a
-> Mixed ('[ 'Just n] ++ MapJust sh) a
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat])
       (sh' :: [Maybe Nat]).
(Rank sh1 ~ Rank sh2) =>
StaticShX sh1
-> StaticShX sh2
-> Proxy sh'
-> Mixed (sh1 ++ sh') a
-> Mixed (sh2 ++ sh') a
forall a (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat])
       (sh' :: [Maybe Nat]).
(Elt a, Rank sh1 ~ Rank sh2) =>
StaticShX sh1
-> StaticShX sh2
-> Proxy sh'
-> Mixed (sh1 ++ sh') a
-> Mixed (sh2 ++ sh') a
mcastPartial (() -> SMayNat () SNat 'Nothing
forall {k} i (f :: k -> *). i -> SMayNat i f 'Nothing
SUnknown () SMayNat () SNat 'Nothing -> StaticShX '[] -> StaticShX '[ 'Nothing]
forall {sh1 :: [Maybe Nat]} (n :: Maybe Nat) (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
SMayNat () SNat n -> StaticShX sh -> StaticShX sh1
:!% StaticShX '[]
forall (sh :: [Maybe Nat]). (sh ~ '[]) => StaticShX sh
ZKX) (SNat n -> SMayNat () SNat ('Just n)
forall {k} (f :: k -> *) (n1 :: k) i.
f n1 -> SMayNat i f ('Just n1)
SKnown SNat n
sn SMayNat () SNat ('Just n) -> StaticShX '[] -> StaticShX '[ 'Just n]
forall {sh1 :: [Maybe Nat]} (n :: Maybe Nat) (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
SMayNat () SNat n -> StaticShX sh -> StaticShX sh1
:!% StaticShX '[]
forall (sh :: [Maybe Nat]). (sh ~ '[]) => StaticShX sh
ZKX) Proxy (MapJust sh)
forall {k} (t :: k). Proxy t
Proxy (Mixed ('[ 'Nothing] ++ MapJust sh) a
 -> Mixed ('[ 'Just n] ++ MapJust sh) a)
-> Mixed ('[ 'Nothing] ++ MapJust sh) a
-> Mixed ('[ 'Just n] ++ MapJust sh) a
forall a b. (a -> b) -> a -> b
$ NonEmpty (Mixed (MapJust sh) a) -> Mixed ('Nothing : MapJust sh) a
forall (sh :: [Maybe Nat]).
NonEmpty (Mixed sh a) -> Mixed ('Nothing : sh) a
forall a (sh :: [Maybe Nat]).
Elt a =>
NonEmpty (Mixed sh a) -> Mixed ('Nothing : sh) a
mfromListOuter (NonEmpty (Shaped sh a) -> NonEmpty (Mixed (MapJust sh) a)
forall a b. Coercible a b => a -> b
coerce NonEmpty (Shaped sh a)
l))

sfromListLinear :: forall sh a. Elt a => ShS sh -> NonEmpty a -> Shaped sh a
sfromListLinear :: forall (sh :: [Nat]) a.
Elt a =>
ShS sh -> NonEmpty a -> Shaped sh a
sfromListLinear ShS sh
sh NonEmpty a
l = Mixed (MapJust sh) a -> Shaped sh a
forall (sh :: [Nat]) a. Mixed (MapJust sh) a -> Shaped sh a
Shaped (IShX (MapJust sh) -> NonEmpty a -> Mixed (MapJust sh) a
forall (sh :: [Maybe Nat]) a.
Elt a =>
IShX sh -> NonEmpty a -> Mixed sh a
mfromListLinear (ShS sh -> IShX (MapJust sh)
forall (sh :: [Nat]). ShS sh -> IShX (MapJust sh)
shxFromShS ShS sh
sh) NonEmpty a
l)

sfromListPrim :: forall n a. PrimElt a => SNat n -> [a] -> Shaped '[n] a
sfromListPrim :: forall (n :: Nat) a. PrimElt a => SNat n -> [a] -> Shaped '[n] a
sfromListPrim SNat n
sn [a]
l
  | ('[ 'Just n] ++ '[]) :~: '[ 'Just n]
Refl <- forall (l :: [Maybe Nat]). (l ++ '[]) :~: l
forall {a} (l :: [a]). (l ++ '[]) :~: l
lemAppNil @'[Just n]
  = let ssh :: StaticShX '[ 'Nothing]
ssh = () -> SMayNat () SNat 'Nothing
forall {k} i (f :: k -> *). i -> SMayNat i f 'Nothing
SUnknown () SMayNat () SNat 'Nothing -> StaticShX '[] -> StaticShX '[ 'Nothing]
forall {sh1 :: [Maybe Nat]} (n :: Maybe Nat) (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
SMayNat () SNat n -> StaticShX sh -> StaticShX sh1
:!% StaticShX '[]
forall (sh :: [Maybe Nat]). (sh ~ '[]) => StaticShX sh
ZKX
        xarr :: XArray ('[ 'Just n] ++ '[]) a
xarr = StaticShX '[ 'Nothing]
-> IShX '[ 'Just n]
-> StaticShX '[]
-> XArray ('[ 'Nothing] ++ '[]) a
-> XArray ('[ 'Just n] ++ '[]) a
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
X.cast StaticShX '[ 'Nothing]
ssh (SNat n -> SMayNat Int SNat ('Just n)
forall {k} (f :: k -> *) (n1 :: k) i.
f n1 -> SMayNat i f ('Just n1)
SKnown SNat n
sn SMayNat Int SNat ('Just n) -> ShX '[] Int -> IShX '[ 'Just n]
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) StaticShX '[]
forall (sh :: [Maybe Nat]). (sh ~ '[]) => StaticShX sh
ZKX (StaticShX '[ 'Nothing] -> [a] -> XArray '[ 'Nothing] a
forall a (n :: Maybe Nat).
Storable a =>
StaticShX '[n] -> [a] -> XArray '[n] a
X.fromList1 StaticShX '[ 'Nothing]
ssh [a]
l)
    in Mixed (MapJust '[n]) a -> Shaped '[n] a
forall (sh :: [Nat]) a. Mixed (MapJust sh) a -> Shaped sh a
Shaped (Mixed (MapJust '[n]) a -> Shaped '[n] a)
-> Mixed (MapJust '[n]) a -> Shaped '[n] a
forall a b. (a -> b) -> a -> b
$ Mixed (MapJust '[n]) (Primitive a) -> Mixed (MapJust '[n]) a
forall (sh :: [Maybe Nat]). Mixed sh (Primitive a) -> Mixed sh a
forall a (sh :: [Maybe Nat]).
PrimElt a =>
Mixed sh (Primitive a) -> Mixed sh a
fromPrimitive (Mixed (MapJust '[n]) (Primitive a) -> Mixed (MapJust '[n]) a)
-> Mixed (MapJust '[n]) (Primitive a) -> Mixed (MapJust '[n]) a
forall a b. (a -> b) -> a -> b
$ IShX '[ 'Just n]
-> XArray '[ 'Just n] a -> Mixed '[ 'Just n] (Primitive a)
forall (sh :: [Maybe Nat]) a.
IShX sh -> XArray sh a -> Mixed sh (Primitive a)
M_Primitive (StaticShX '[ 'Just n] -> XArray '[ 'Just n] a -> IShX '[ 'Just n]
forall (sh :: [Maybe Nat]) a.
StaticShX sh -> XArray sh a -> IShX sh
X.shape (SNat n -> SMayNat () SNat ('Just n)
forall {k} (f :: k -> *) (n1 :: k) i.
f n1 -> SMayNat i f ('Just n1)
SKnown SNat n
sn SMayNat () SNat ('Just n) -> StaticShX '[] -> StaticShX '[ 'Just n]
forall {sh1 :: [Maybe Nat]} (n :: Maybe Nat) (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
SMayNat () SNat n -> StaticShX sh -> StaticShX sh1
:!% StaticShX '[]
forall (sh :: [Maybe Nat]). (sh ~ '[]) => StaticShX sh
ZKX) XArray '[ 'Just n] a
XArray ('[ 'Just n] ++ '[]) a
xarr) XArray '[ 'Just n] a
XArray ('[ 'Just n] ++ '[]) a
xarr

sfromListPrimLinear :: PrimElt a => ShS sh -> [a] -> Shaped sh a
sfromListPrimLinear :: forall a (sh :: [Nat]). PrimElt a => ShS sh -> [a] -> Shaped sh a
sfromListPrimLinear ShS sh
sh [a]
l =
  let M_Primitive IShX '[ 'Nothing]
_ XArray '[ 'Nothing] a
xarr = Mixed '[ 'Nothing] a -> Mixed '[ 'Nothing] (Primitive a)
forall (sh :: [Maybe Nat]). Mixed sh a -> Mixed sh (Primitive a)
forall a (sh :: [Maybe Nat]).
PrimElt a =>
Mixed sh a -> Mixed sh (Primitive a)
toPrimitive ([a] -> Mixed '[ 'Nothing] a
forall a. PrimElt a => [a] -> Mixed '[ 'Nothing] a
mfromListPrim [a]
l)
  in Mixed (MapJust sh) a -> Shaped sh a
forall (sh :: [Nat]) a. Mixed (MapJust sh) a -> Shaped sh a
Shaped (Mixed (MapJust sh) a -> Shaped sh a)
-> Mixed (MapJust sh) a -> Shaped sh a
forall a b. (a -> b) -> a -> b
$ Mixed (MapJust sh) (Primitive a) -> Mixed (MapJust sh) a
forall (sh :: [Maybe Nat]). Mixed sh (Primitive a) -> Mixed sh a
forall a (sh :: [Maybe Nat]).
PrimElt a =>
Mixed sh (Primitive a) -> Mixed sh a
fromPrimitive (Mixed (MapJust sh) (Primitive a) -> Mixed (MapJust sh) a)
-> Mixed (MapJust sh) (Primitive a) -> Mixed (MapJust sh) a
forall a b. (a -> b) -> a -> b
$ IShX (MapJust sh)
-> XArray (MapJust sh) a -> Mixed (MapJust sh) (Primitive a)
forall (sh :: [Maybe Nat]) a.
IShX sh -> XArray sh a -> Mixed sh (Primitive a)
M_Primitive (ShS sh -> IShX (MapJust sh)
forall (sh :: [Nat]). ShS sh -> IShX (MapJust sh)
shxFromShS ShS sh
sh) (StaticShX '[ 'Nothing]
-> IShX (MapJust sh)
-> XArray '[ 'Nothing] a
-> XArray (MapJust sh) a
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
Storable a =>
StaticShX sh1 -> IShX sh2 -> XArray sh1 a -> XArray sh2 a
X.reshape (() -> SMayNat () SNat 'Nothing
forall {k} i (f :: k -> *). i -> SMayNat i f 'Nothing
SUnknown () SMayNat () SNat 'Nothing -> StaticShX '[] -> StaticShX '[ 'Nothing]
forall {sh1 :: [Maybe Nat]} (n :: Maybe Nat) (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
SMayNat () SNat n -> StaticShX sh -> StaticShX sh1
:!% StaticShX '[]
forall (sh :: [Maybe Nat]). (sh ~ '[]) => StaticShX sh
ZKX) (ShS sh -> IShX (MapJust sh)
forall (sh :: [Nat]). ShS sh -> IShX (MapJust sh)
shxFromShS ShS sh
sh) XArray '[ 'Nothing] a
xarr)

stoList :: Elt a => Shaped '[n] a -> [a]
stoList :: forall a (n :: Nat). Elt a => Shaped '[n] a -> [a]
stoList = (Shaped '[] a -> a) -> [Shaped '[] a] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map Shaped '[] a -> a
forall a. Elt a => Shaped '[] a -> a
sunScalar ([Shaped '[] a] -> [a])
-> (Shaped '[n] a -> [Shaped '[] a]) -> Shaped '[n] a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shaped '[n] a -> [Shaped '[] a]
forall a (n :: Nat) (sh :: [Nat]).
Elt a =>
Shaped (n : sh) a -> [Shaped sh a]
stoListOuter

stoListOuter :: Elt a => Shaped (n : sh) a -> [Shaped sh a]
stoListOuter :: forall a (n :: Nat) (sh :: [Nat]).
Elt a =>
Shaped (n : sh) a -> [Shaped sh a]
stoListOuter (Shaped Mixed (MapJust (n : sh)) a
arr) = [Mixed (MapJust sh) a] -> [Shaped sh a]
forall a b. Coercible a b => a -> b
coerce (Mixed ('Just n : MapJust sh) a -> [Mixed (MapJust sh) a]
forall (n :: Maybe Nat) (sh :: [Maybe Nat]).
Mixed (n : sh) a -> [Mixed sh a]
forall a (n :: Maybe Nat) (sh :: [Maybe Nat]).
Elt a =>
Mixed (n : sh) a -> [Mixed sh a]
mtoListOuter Mixed ('Just n : MapJust sh) a
Mixed (MapJust (n : sh)) a
arr)

stoListLinear :: Elt a => Shaped sh a -> [a]
stoListLinear :: forall a (sh :: [Nat]). Elt a => Shaped sh a -> [a]
stoListLinear (Shaped Mixed (MapJust sh) a
arr) = Mixed (MapJust sh) a -> [a]
forall a (sh :: [Maybe Nat]). Elt a => Mixed sh a -> [a]
mtoListLinear Mixed (MapJust sh) a
arr

sfromOrthotope :: PrimElt a => ShS sh -> SS.Array sh a -> Shaped sh a
sfromOrthotope :: forall a (sh :: [Nat]).
PrimElt a =>
ShS sh -> Array sh a -> Shaped sh a
sfromOrthotope ShS sh
sh (SS.A (SG.A T Vector a
arr)) =
  Mixed (MapJust sh) a -> Shaped sh a
forall (sh :: [Nat]) a. Mixed (MapJust sh) a -> Shaped sh a
Shaped (Mixed (MapJust sh) (Primitive a) -> Mixed (MapJust sh) a
forall (sh :: [Maybe Nat]). Mixed sh (Primitive a) -> Mixed sh a
forall a (sh :: [Maybe Nat]).
PrimElt a =>
Mixed sh (Primitive a) -> Mixed sh a
fromPrimitive (IShX (MapJust sh)
-> XArray (MapJust sh) a -> Mixed (MapJust sh) (Primitive a)
forall (sh :: [Maybe Nat]) a.
IShX sh -> XArray sh a -> Mixed sh (Primitive a)
M_Primitive (ShS sh -> IShX (MapJust sh)
forall (sh :: [Nat]). ShS sh -> IShX (MapJust sh)
shxFromShS ShS sh
sh) (Array (Rank (MapJust sh)) a -> XArray (MapJust sh) a
forall (sh :: [Maybe Nat]) a. Array (Rank sh) a -> XArray sh a
X.XArray (Array (Rank (MapJust sh)) Vector a -> Array (Rank (MapJust sh)) a
forall (n :: Nat) a. Array n Vector a -> Array n a
RS.A (ShapeL -> T Vector a -> Array (Rank (MapJust sh)) Vector a
forall (n :: Nat) (v :: * -> *) a. ShapeL -> T v a -> Array n v a
RG.A (ShS sh -> ShapeL
forall (sh :: [Nat]). ShS sh -> ShapeL
shsToList ShS sh
sh) T Vector a
arr)))))

stoOrthotope :: PrimElt a => Shaped sh a -> SS.Array sh a
stoOrthotope :: forall a (sh :: [Nat]). PrimElt a => Shaped sh a -> Array sh a
stoOrthotope (Shaped sh a -> Shaped sh (Primitive a)
forall a (sh :: [Nat]).
PrimElt a =>
Shaped sh a -> Shaped sh (Primitive a)
stoPrimitive -> Shaped (M_Primitive IShX (MapJust sh)
_ (X.XArray (RS.A (RG.A ShapeL
_ T Vector a
arr))))) = Array sh Vector a -> Array sh a
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
SS.A (T Vector a -> Array sh Vector a
forall (sh :: [Nat]) (v :: * -> *) a. T v a -> Array sh v a
SG.A T Vector a
arr)

sunScalar :: Elt a => Shaped '[] a -> a
sunScalar :: forall a. Elt a => Shaped '[] a -> a
sunScalar Shaped '[] a
arr = Shaped '[] a -> IIxS '[] -> a
forall a (sh :: [Nat]). Elt a => Shaped sh a -> IIxS sh -> a
sindex Shaped '[] a
arr IIxS '[]
forall (sh :: [Nat]) i. (sh ~ '[]) => IxS sh i
ZIS

snest :: forall sh sh' a. Elt a => ShS sh -> Shaped (sh ++ sh') a -> Shaped sh (Shaped sh' a)
snest :: forall (sh :: [Nat]) (sh' :: [Nat]) a.
Elt a =>
ShS sh -> Shaped (sh ++ sh') a -> Shaped sh (Shaped sh' a)
snest ShS sh
sh Shaped (sh ++ sh') a
arr
  | MapJust (sh ++ sh') :~: (MapJust sh ++ MapJust sh')
Refl <- ShS sh
-> Proxy sh' -> MapJust (sh ++ sh') :~: (MapJust sh ++ MapJust sh')
forall (sh1 :: [Nat]) (sh2 :: [Nat]).
ShS sh1
-> Proxy sh2
-> MapJust (sh1 ++ sh2) :~: (MapJust sh1 ++ MapJust sh2)
lemMapJustApp ShS sh
sh (forall (t :: [Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @sh')
  = Mixed (MapJust sh) (Mixed (MapJust sh') a)
-> Shaped sh (Shaped sh' a)
forall a b. Coercible a b => a -> b
coerce (StaticShX (MapJust sh)
-> Mixed (MapJust sh ++ MapJust sh') a
-> Mixed (MapJust sh) (Mixed (MapJust sh') a)
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) a.
Elt a =>
StaticShX sh -> Mixed (sh ++ sh') a -> Mixed sh (Mixed sh' a)
mnest (ShX (MapJust sh) Int -> StaticShX (MapJust sh)
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX (ShS sh -> ShX (MapJust sh) Int
forall (sh :: [Nat]). ShS sh -> IShX (MapJust sh)
shxFromShS ShS sh
sh)) (Shaped (sh ++ sh') a -> Mixed (MapJust (sh ++ sh')) a
forall a b. Coercible a b => a -> b
coerce Shaped (sh ++ sh') a
arr))

sunNest :: forall sh sh' a. Elt a => Shaped sh (Shaped sh' a) -> Shaped (sh ++ sh') a
sunNest :: forall (sh :: [Nat]) (sh' :: [Nat]) a.
Elt a =>
Shaped sh (Shaped sh' a) -> Shaped (sh ++ sh') a
sunNest sarr :: Shaped sh (Shaped sh' a)
sarr@(Shaped (M_Shaped (M_Nest IShX (MapJust sh)
_ Mixed (MapJust sh ++ MapJust sh') a
arr)))
  | MapJust (sh ++ sh') :~: (MapJust sh ++ MapJust sh')
Refl <- ShS sh
-> Proxy sh' -> MapJust (sh ++ sh') :~: (MapJust sh ++ MapJust sh')
forall (sh1 :: [Nat]) (sh2 :: [Nat]).
ShS sh1
-> Proxy sh2
-> MapJust (sh1 ++ sh2) :~: (MapJust sh1 ++ MapJust sh2)
lemMapJustApp (Shaped sh (Shaped sh' a) -> ShS sh
forall (sh :: [Nat]) a. Elt a => Shaped sh a -> ShS sh
sshape Shaped sh (Shaped sh' a)
sarr) (forall (t :: [Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @sh')
  = Mixed (MapJust (sh ++ sh')) a -> Shaped (sh ++ sh') a
forall (sh :: [Nat]) a. Mixed (MapJust sh) a -> Shaped sh a
Shaped Mixed (MapJust (sh ++ sh')) a
Mixed (MapJust sh ++ MapJust sh') a
arr

szip :: (Elt a, Elt b) => Shaped sh a -> Shaped sh b -> Shaped sh (a, b)
szip :: forall a b (sh :: [Nat]).
(Elt a, Elt b) =>
Shaped sh a -> Shaped sh b -> Shaped sh (a, b)
szip = (Mixed (MapJust sh) a
 -> Mixed (MapJust sh) b -> Mixed (MapJust sh) (a, b))
-> Shaped sh a -> Shaped sh b -> Shaped sh (a, b)
forall a b. Coercible a b => a -> b
coerce Mixed (MapJust sh) a
-> Mixed (MapJust sh) b -> Mixed (MapJust sh) (a, b)
forall a b (sh :: [Maybe Nat]).
(Elt a, Elt b) =>
Mixed sh a -> Mixed sh b -> Mixed sh (a, b)
mzip

sunzip :: Shaped sh (a, b) -> (Shaped sh a, Shaped sh b)
sunzip :: forall (sh :: [Nat]) a b.
Shaped sh (a, b) -> (Shaped sh a, Shaped sh b)
sunzip = (Mixed (MapJust sh) (a, b)
 -> (Mixed (MapJust sh) a, Mixed (MapJust sh) b))
-> Shaped sh (a, b) -> (Shaped sh a, Shaped sh b)
forall a b. Coercible a b => a -> b
coerce Mixed (MapJust sh) (a, b)
-> (Mixed (MapJust sh) a, Mixed (MapJust sh) b)
forall (sh :: [Maybe Nat]) a b.
Mixed sh (a, b) -> (Mixed sh a, Mixed sh b)
munzip

srerankP :: forall sh1 sh2 sh a b. (Storable a, Storable b)
         => ShS sh -> ShS sh2
         -> (Shaped sh1 (Primitive a) -> Shaped sh2 (Primitive b))
         -> Shaped (sh ++ sh1) (Primitive a) -> Shaped (sh ++ sh2) (Primitive b)
srerankP :: forall (sh1 :: [Nat]) (sh2 :: [Nat]) (sh :: [Nat]) a b.
(Storable a, Storable b) =>
ShS sh
-> ShS sh2
-> (Shaped sh1 (Primitive a) -> Shaped sh2 (Primitive b))
-> Shaped (sh ++ sh1) (Primitive a)
-> Shaped (sh ++ sh2) (Primitive b)
srerankP ShS sh
sh ShS sh2
sh2 Shaped sh1 (Primitive a) -> Shaped sh2 (Primitive b)
f sarr :: Shaped (sh ++ sh1) (Primitive a)
sarr@(Shaped Mixed (MapJust (sh ++ sh1)) (Primitive a)
arr)
  | MapJust (sh ++ sh1) :~: (MapJust sh ++ MapJust sh1)
Refl <- ShS sh
-> Proxy sh1 -> MapJust (sh ++ sh1) :~: (MapJust sh ++ MapJust sh1)
forall (sh1 :: [Nat]) (sh2 :: [Nat]).
ShS sh1
-> Proxy sh2
-> MapJust (sh1 ++ sh2) :~: (MapJust sh1 ++ MapJust sh2)
lemMapJustApp ShS sh
sh (forall (t :: [Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @sh1)
  , MapJust (sh ++ sh2) :~: (MapJust sh ++ MapJust sh2)
Refl <- ShS sh
-> Proxy sh2 -> MapJust (sh ++ sh2) :~: (MapJust sh ++ MapJust sh2)
forall (sh1 :: [Nat]) (sh2 :: [Nat]).
ShS sh1
-> Proxy sh2
-> MapJust (sh1 ++ sh2) :~: (MapJust sh1 ++ MapJust sh2)
lemMapJustApp ShS sh
sh (forall (t :: [Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @sh2)
  = Mixed (MapJust (sh ++ sh2)) (Primitive b)
-> Shaped (sh ++ sh2) (Primitive b)
forall (sh :: [Nat]) a. Mixed (MapJust sh) a -> Shaped sh a
Shaped (StaticShX (MapJust sh)
-> IShX (MapJust sh2)
-> (Mixed (MapJust sh1) (Primitive a)
    -> Mixed (MapJust sh2) (Primitive b))
-> Mixed (MapJust sh ++ MapJust sh1) (Primitive a)
-> Mixed (MapJust sh ++ MapJust sh2) (Primitive b)
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat])
       (sh :: [Maybe Nat]) a b.
(Storable a, Storable b) =>
StaticShX sh
-> IShX sh2
-> (Mixed sh1 (Primitive a) -> Mixed sh2 (Primitive b))
-> Mixed (sh ++ sh1) (Primitive a)
-> Mixed (sh ++ sh2) (Primitive b)
mrerankP (ShX (MapJust sh) Int -> StaticShX (MapJust sh)
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX (Proxy (MapJust sh1)
-> StaticShX (MapJust sh)
-> ShX (MapJust sh ++ MapJust sh1) Int
-> ShX (MapJust sh) Int
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) i
       (proxy :: [Maybe Nat] -> *).
proxy sh' -> StaticShX sh -> ShX (sh ++ sh') i -> ShX sh i
shxTakeSSX (forall (t :: [Maybe Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @(MapJust sh1)) (ShX (MapJust sh) Int -> StaticShX (MapJust sh)
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX (ShS sh -> ShX (MapJust sh) Int
forall (sh :: [Nat]). ShS sh -> IShX (MapJust sh)
shxFromShS ShS sh
sh)) (ShS (sh ++ sh1) -> IShX (MapJust (sh ++ sh1))
forall (sh :: [Nat]). ShS sh -> IShX (MapJust sh)
shxFromShS (Shaped (sh ++ sh1) (Primitive a) -> ShS (sh ++ sh1)
forall (sh :: [Nat]) a. Elt a => Shaped sh a -> ShS sh
sshape Shaped (sh ++ sh1) (Primitive a)
sarr))))
                     (ShS sh2 -> IShX (MapJust sh2)
forall (sh :: [Nat]). ShS sh -> IShX (MapJust sh)
shxFromShS ShS sh2
sh2)
                     (\Mixed (MapJust sh1) (Primitive a)
a -> let Shaped Mixed (MapJust sh2) (Primitive b)
r = Shaped sh1 (Primitive a) -> Shaped sh2 (Primitive b)
f (Mixed (MapJust sh1) (Primitive a) -> Shaped sh1 (Primitive a)
forall (sh :: [Nat]) a. Mixed (MapJust sh) a -> Shaped sh a
Shaped Mixed (MapJust sh1) (Primitive a)
a) in Mixed (MapJust sh2) (Primitive b)
r)
                     Mixed (MapJust (sh ++ sh1)) (Primitive a)
Mixed (MapJust sh ++ MapJust sh1) (Primitive a)
arr)

srerank :: forall sh1 sh2 sh a b. (PrimElt a, PrimElt b)
        => ShS sh -> ShS sh2
        -> (Shaped sh1 a -> Shaped sh2 b)
        -> Shaped (sh ++ sh1) a -> Shaped (sh ++ sh2) b
srerank :: forall (sh1 :: [Nat]) (sh2 :: [Nat]) (sh :: [Nat]) a b.
(PrimElt a, PrimElt b) =>
ShS sh
-> ShS sh2
-> (Shaped sh1 a -> Shaped sh2 b)
-> Shaped (sh ++ sh1) a
-> Shaped (sh ++ sh2) b
srerank ShS sh
sh ShS sh2
sh2 Shaped sh1 a -> Shaped sh2 b
f (Shaped (sh ++ sh1) a -> Shaped (sh ++ sh1) (Primitive a)
forall a (sh :: [Nat]).
PrimElt a =>
Shaped sh a -> Shaped sh (Primitive a)
stoPrimitive -> Shaped (sh ++ sh1) (Primitive a)
arr) =
  Shaped (sh ++ sh2) (Primitive b) -> Shaped (sh ++ sh2) b
forall a (sh :: [Nat]).
PrimElt a =>
Shaped sh (Primitive a) -> Shaped sh a
sfromPrimitive (Shaped (sh ++ sh2) (Primitive b) -> Shaped (sh ++ sh2) b)
-> Shaped (sh ++ sh2) (Primitive b) -> Shaped (sh ++ sh2) b
forall a b. (a -> b) -> a -> b
$ ShS sh
-> ShS sh2
-> (Shaped sh1 (Primitive a) -> Shaped sh2 (Primitive b))
-> Shaped (sh ++ sh1) (Primitive a)
-> Shaped (sh ++ sh2) (Primitive b)
forall (sh1 :: [Nat]) (sh2 :: [Nat]) (sh :: [Nat]) a b.
(Storable a, Storable b) =>
ShS sh
-> ShS sh2
-> (Shaped sh1 (Primitive a) -> Shaped sh2 (Primitive b))
-> Shaped (sh ++ sh1) (Primitive a)
-> Shaped (sh ++ sh2) (Primitive b)
srerankP ShS sh
sh ShS sh2
sh2 (Shaped sh2 b -> Shaped sh2 (Primitive b)
forall a (sh :: [Nat]).
PrimElt a =>
Shaped sh a -> Shaped sh (Primitive a)
stoPrimitive (Shaped sh2 b -> Shaped sh2 (Primitive b))
-> (Shaped sh1 (Primitive a) -> Shaped sh2 b)
-> Shaped sh1 (Primitive a)
-> Shaped sh2 (Primitive b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shaped sh1 a -> Shaped sh2 b
f (Shaped sh1 a -> Shaped sh2 b)
-> (Shaped sh1 (Primitive a) -> Shaped sh1 a)
-> Shaped sh1 (Primitive a)
-> Shaped sh2 b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Shaped sh1 (Primitive a) -> Shaped sh1 a
forall a (sh :: [Nat]).
PrimElt a =>
Shaped sh (Primitive a) -> Shaped sh a
sfromPrimitive) Shaped (sh ++ sh1) (Primitive a)
arr

sreplicate :: forall sh sh' a. Elt a => ShS sh -> Shaped sh' a -> Shaped (sh ++ sh') a
sreplicate :: forall (sh :: [Nat]) (sh' :: [Nat]) a.
Elt a =>
ShS sh -> Shaped sh' a -> Shaped (sh ++ sh') a
sreplicate ShS sh
sh (Shaped Mixed (MapJust sh') a
arr)
  | MapJust (sh ++ sh') :~: (MapJust sh ++ MapJust sh')
Refl <- ShS sh
-> Proxy sh' -> MapJust (sh ++ sh') :~: (MapJust sh ++ MapJust sh')
forall (sh1 :: [Nat]) (sh2 :: [Nat]).
ShS sh1
-> Proxy sh2
-> MapJust (sh1 ++ sh2) :~: (MapJust sh1 ++ MapJust sh2)
lemMapJustApp ShS sh
sh (forall (t :: [Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @sh')
  = Mixed (MapJust (sh ++ sh')) a -> Shaped (sh ++ sh') a
forall (sh :: [Nat]) a. Mixed (MapJust sh) a -> Shaped sh a
Shaped (IShX (MapJust sh)
-> Mixed (MapJust sh') a -> Mixed (MapJust sh ++ MapJust sh') a
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) a.
Elt a =>
IShX sh -> Mixed sh' a -> Mixed (sh ++ sh') a
mreplicate (ShS sh -> IShX (MapJust sh)
forall (sh :: [Nat]). ShS sh -> IShX (MapJust sh)
shxFromShS ShS sh
sh) Mixed (MapJust sh') a
arr)

sreplicateScalP :: forall sh a. Storable a => ShS sh -> a -> Shaped sh (Primitive a)
sreplicateScalP :: forall (sh :: [Nat]) a.
Storable a =>
ShS sh -> a -> Shaped sh (Primitive a)
sreplicateScalP ShS sh
sh a
x = Mixed (MapJust sh) (Primitive a) -> Shaped sh (Primitive a)
forall (sh :: [Nat]) a. Mixed (MapJust sh) a -> Shaped sh a
Shaped (IShX (MapJust sh) -> a -> Mixed (MapJust sh) (Primitive a)
forall (sh :: [Maybe Nat]) a.
Storable a =>
IShX sh -> a -> Mixed sh (Primitive a)
mreplicateScalP (ShS sh -> IShX (MapJust sh)
forall (sh :: [Nat]). ShS sh -> IShX (MapJust sh)
shxFromShS ShS sh
sh) a
x)

sreplicateScal :: PrimElt a => ShS sh -> a -> Shaped sh a
sreplicateScal :: forall a (sh :: [Nat]). PrimElt a => ShS sh -> a -> Shaped sh a
sreplicateScal ShS sh
sh a
x = Shaped sh (Primitive a) -> Shaped sh a
forall a (sh :: [Nat]).
PrimElt a =>
Shaped sh (Primitive a) -> Shaped sh a
sfromPrimitive (ShS sh -> a -> Shaped sh (Primitive a)
forall (sh :: [Nat]) a.
Storable a =>
ShS sh -> a -> Shaped sh (Primitive a)
sreplicateScalP ShS sh
sh a
x)

sslice :: Elt a => SNat i -> SNat n -> Shaped (i + n + k : sh) a -> Shaped (n : sh) a
sslice :: forall a (i :: Nat) (n :: Nat) (k :: Nat) (sh :: [Nat]).
Elt a =>
SNat i
-> SNat n -> Shaped (((i + n) + k) : sh) a -> Shaped (n : sh) a
sslice SNat i
i n :: SNat n
n@SNat n
SNat Shaped (((i + n) + k) : sh) a
arr =
  let SNat n
_ :$$ ShS sh
ShS sh
sh = Shaped (((i + n) + k) : sh) a -> ShS (((i + n) + k) : sh)
forall (sh :: [Nat]) a. Elt a => Shaped sh a -> ShS sh
sshape Shaped (((i + n) + k) : sh) a
arr
  in ShS (n : sh)
-> (forall (sh' :: [Maybe Nat]) b.
    Storable b =>
    StaticShX sh'
    -> XArray (MapJust (((i + n) + k) : sh) ++ sh') b
    -> XArray (MapJust (n : sh) ++ sh') b)
-> Shaped (((i + n) + k) : sh) a
-> Shaped (n : sh) a
forall (sh1 :: [Nat]) (sh2 :: [Nat]) a.
Elt a =>
ShS sh2
-> (forall (sh' :: [Maybe Nat]) b.
    Storable b =>
    StaticShX sh'
    -> XArray (MapJust sh1 ++ sh') b -> XArray (MapJust sh2 ++ sh') b)
-> Shaped sh1 a
-> Shaped sh2 a
slift (SNat n
n SNat n -> ShS sh -> ShS (n : sh)
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, (n : sh) ~ sh1) =>
SNat n -> ShS sh -> ShS sh1
:$$ ShS sh
sh) (\StaticShX sh'
_ -> SNat i
-> SNat n
-> XArray ('Just ((i + n) + k) : (MapJust sh ++ sh')) b
-> XArray ('Just n : (MapJust sh ++ sh')) b
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
X.slice SNat i
i SNat n
n) Shaped (((i + n) + k) : sh) a
arr

srev1 :: Elt a => Shaped (n : sh) a -> Shaped (n : sh) a
srev1 :: forall a (n :: Nat) (sh :: [Nat]).
Elt a =>
Shaped (n : sh) a -> Shaped (n : sh) a
srev1 Shaped (n : sh) a
arr = ShS (n : sh)
-> (forall (sh' :: [Maybe Nat]) b.
    Storable b =>
    StaticShX sh'
    -> XArray (MapJust (n : sh) ++ sh') b
    -> XArray (MapJust (n : sh) ++ sh') b)
-> Shaped (n : sh) a
-> Shaped (n : sh) a
forall (sh1 :: [Nat]) (sh2 :: [Nat]) a.
Elt a =>
ShS sh2
-> (forall (sh' :: [Maybe Nat]) b.
    Storable b =>
    StaticShX sh'
    -> XArray (MapJust sh1 ++ sh') b -> XArray (MapJust sh2 ++ sh') b)
-> Shaped sh1 a
-> Shaped sh2 a
slift (Shaped (n : sh) a -> ShS (n : sh)
forall (sh :: [Nat]) a. Elt a => Shaped sh a -> ShS sh
sshape Shaped (n : sh) a
arr) (\StaticShX sh'
_ -> XArray ('Just n : (MapJust sh ++ sh')) b
-> XArray ('Just n : (MapJust sh ++ sh')) b
XArray (MapJust (n : sh) ++ sh') b
-> XArray (MapJust (n : sh) ++ sh') b
forall (n :: Maybe Nat) (sh :: [Maybe Nat]) a.
XArray (n : sh) a -> XArray (n : sh) a
X.rev1) Shaped (n : sh) a
arr

sreshape :: (Elt a, Product sh ~ Product sh') => ShS sh' -> Shaped sh a -> Shaped sh' a
sreshape :: forall a (sh :: [Nat]) (sh' :: [Nat]).
(Elt a, Product sh ~ Product sh') =>
ShS sh' -> Shaped sh a -> Shaped sh' a
sreshape ShS sh'
sh' (Shaped Mixed (MapJust sh) a
arr) = Mixed (MapJust sh') a -> Shaped sh' a
forall (sh :: [Nat]) a. Mixed (MapJust sh) a -> Shaped sh a
Shaped (IShX (MapJust sh') -> Mixed (MapJust sh) a -> Mixed (MapJust sh') a
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) a.
Elt a =>
IShX sh' -> Mixed sh a -> Mixed sh' a
mreshape (ShS sh' -> IShX (MapJust sh')
forall (sh :: [Nat]). ShS sh -> IShX (MapJust sh)
shxFromShS ShS sh'
sh') Mixed (MapJust sh) a
arr)

sflatten :: Elt a => Shaped sh a -> Shaped '[Product sh] a
sflatten :: forall a (sh :: [Nat]).
Elt a =>
Shaped sh a -> Shaped '[Product sh] a
sflatten Shaped sh a
arr =
  case ShS sh -> SNat (Product sh)
forall (sh :: [Nat]). ShS sh -> SNat (Product sh)
shsProduct (Shaped sh a -> ShS sh
forall (sh :: [Nat]) a. Elt a => Shaped sh a -> ShS sh
sshape Shaped sh a
arr) of  -- TODO: simplify when removing the KnownNat stuff
    n :: SNat (Product sh)
n@SNat (Product sh)
SNat -> ShS '[Product sh] -> Shaped sh a -> Shaped '[Product sh] a
forall a (sh :: [Nat]) (sh' :: [Nat]).
(Elt a, Product sh ~ Product sh') =>
ShS sh' -> Shaped sh a -> Shaped sh' a
sreshape (SNat (Product sh)
n SNat (Product sh) -> ShS '[] -> ShS '[Product sh]
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, (n : sh) ~ sh1) =>
SNat n -> ShS sh -> ShS sh1
:$$ ShS '[]
forall (sh :: [Nat]). (sh ~ '[]) => ShS sh
ZSS) Shaped sh a
arr

siota :: (Enum a, PrimElt a) => SNat n -> Shaped '[n] a
siota :: forall a (n :: Nat). (Enum a, PrimElt a) => SNat n -> Shaped '[n] a
siota SNat n
sn = Mixed (MapJust '[n]) a -> Shaped '[n] a
forall (sh :: [Nat]) a. Mixed (MapJust sh) a -> Shaped sh a
Shaped (SNat n -> Mixed '[ 'Just n] a
forall a (n :: Nat).
(Enum a, PrimElt a) =>
SNat n -> Mixed '[ 'Just n] a
miota SNat n
sn)

-- | Throws if the array is empty.
sminIndexPrim :: (PrimElt a, NumElt a) => Shaped sh a -> IIxS sh
sminIndexPrim :: forall a (sh :: [Nat]).
(PrimElt a, NumElt a) =>
Shaped sh a -> IIxS sh
sminIndexPrim sarr :: Shaped sh a
sarr@(Shaped Mixed (MapJust sh) a
arr) = ShS sh -> IxX (MapJust sh) Int -> IxS sh Int
forall (sh :: [Nat]) i. ShS sh -> IxX (MapJust sh) i -> IxS sh i
ixsFromIxX (Shaped sh (Primitive a) -> ShS sh
forall (sh :: [Nat]) a. Elt a => Shaped sh a -> ShS sh
sshape (Shaped sh a -> Shaped sh (Primitive a)
forall a (sh :: [Nat]).
PrimElt a =>
Shaped sh a -> Shaped sh (Primitive a)
stoPrimitive Shaped sh a
sarr)) (Mixed (MapJust sh) a -> IxX (MapJust sh) Int
forall a (sh :: [Maybe Nat]).
(PrimElt a, NumElt a) =>
Mixed sh a -> IIxX sh
mminIndexPrim Mixed (MapJust sh) a
arr)

-- | Throws if the array is empty.
smaxIndexPrim :: (PrimElt a, NumElt a) => Shaped sh a -> IIxS sh
smaxIndexPrim :: forall a (sh :: [Nat]).
(PrimElt a, NumElt a) =>
Shaped sh a -> IIxS sh
smaxIndexPrim sarr :: Shaped sh a
sarr@(Shaped Mixed (MapJust sh) a
arr) = ShS sh -> IxX (MapJust sh) Int -> IxS sh Int
forall (sh :: [Nat]) i. ShS sh -> IxX (MapJust sh) i -> IxS sh i
ixsFromIxX (Shaped sh (Primitive a) -> ShS sh
forall (sh :: [Nat]) a. Elt a => Shaped sh a -> ShS sh
sshape (Shaped sh a -> Shaped sh (Primitive a)
forall a (sh :: [Nat]).
PrimElt a =>
Shaped sh a -> Shaped sh (Primitive a)
stoPrimitive Shaped sh a
sarr)) (Mixed (MapJust sh) a -> IxX (MapJust sh) Int
forall a (sh :: [Maybe Nat]).
(PrimElt a, NumElt a) =>
Mixed sh a -> IIxX sh
mmaxIndexPrim Mixed (MapJust sh) a
arr)

sdot1Inner :: forall sh n a. (PrimElt a, NumElt a)
           => Proxy n -> Shaped (sh ++ '[n]) a -> Shaped (sh ++ '[n]) a -> Shaped sh a
sdot1Inner :: forall (sh :: [Nat]) (n :: Nat) a.
(PrimElt a, NumElt a) =>
Proxy n
-> Shaped (sh ++ '[n]) a -> Shaped (sh ++ '[n]) a -> Shaped sh a
sdot1Inner Proxy n
Proxy sarr1 :: Shaped (sh ++ '[n]) a
sarr1@(Shaped Mixed (MapJust (sh ++ '[n])) a
arr1) (Shaped Mixed (MapJust (sh ++ '[n])) a
arr2)
  | Init (sh ++ '[n]) :~: sh
Refl <- Proxy sh -> Proxy n -> Init (sh ++ '[n]) :~: sh
forall {a} (l :: [a]) (x :: a).
Proxy l -> Proxy x -> Init (l ++ '[x]) :~: l
lemInitApp (forall (t :: [Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @sh) (forall (t :: Nat). Proxy t
forall {k} (t :: k). Proxy t
Proxy @n)
  , Last (sh ++ '[n]) :~: n
Refl <- Proxy sh -> Proxy n -> Last (sh ++ '[n]) :~: n
forall {k} (l :: [k]) (x :: k).
Proxy l -> Proxy x -> Last (l ++ '[x]) :~: x
lemLastApp (forall (t :: [Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @sh) (forall (t :: Nat). Proxy t
forall {k} (t :: k). Proxy t
Proxy @n)
  = case Shaped (sh ++ '[n]) a -> ShS (sh ++ '[n])
forall (sh :: [Nat]) a. Elt a => Shaped sh a -> ShS sh
sshape Shaped (sh ++ '[n]) a
sarr1 of
      SNat n
_ :$$ ShS sh
_
        | MapJust (Init (n : sh) ++ '[n])
:~: (MapJust (Init (n : sh)) ++ MapJust '[n])
Refl <- ShS (Init (n : sh))
-> Proxy '[n]
-> MapJust (Init (n : sh) ++ '[n])
   :~: (MapJust (Init (n : sh)) ++ MapJust '[n])
forall (sh1 :: [Nat]) (sh2 :: [Nat]).
ShS sh1
-> Proxy sh2
-> MapJust (sh1 ++ sh2) :~: (MapJust sh1 ++ MapJust sh2)
lemMapJustApp (ShS (n : sh) -> ShS (Init (n : sh))
forall (n :: Nat) (sh :: [Nat]).
ShS (n : sh) -> ShS (Init (n : sh))
shsInit (Shaped (n : sh) a -> ShS (n : sh)
forall (sh :: [Nat]) a. Elt a => Shaped sh a -> ShS sh
sshape Shaped (n : sh) a
Shaped (sh ++ '[n]) a
sarr1)) (forall (t :: [Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @'[n])
        -> Mixed (MapJust sh) a -> Shaped sh a
forall (sh :: [Nat]) a. Mixed (MapJust sh) a -> Shaped sh a
Shaped (Proxy ('Just n)
-> Mixed (MapJust sh ++ '[ 'Just n]) a
-> Mixed (MapJust sh ++ '[ 'Just n]) a
-> Mixed (MapJust sh) a
forall (sh :: [Maybe Nat]) (n :: Maybe Nat) a.
(PrimElt a, NumElt a) =>
Proxy n
-> Mixed (sh ++ '[n]) a -> Mixed (sh ++ '[n]) a -> Mixed sh a
mdot1Inner (forall (t :: Maybe Nat). Proxy t
forall {k} (t :: k). Proxy t
Proxy @(Just n)) Mixed (MapJust (sh ++ '[n])) a
Mixed (MapJust sh ++ '[ 'Just n]) a
arr1 Mixed (MapJust (sh ++ '[n])) a
Mixed (MapJust sh ++ '[ 'Just n]) a
arr2)
      ShS (sh ++ '[n])
_ -> [Char] -> Shaped sh a
forall a. HasCallStack => [Char] -> a
error [Char]
"unreachable"

-- | This has a temporary, suboptimal implementation in terms of 'mflatten'.
-- Prefer 'sdot1Inner' if applicable.
sdot :: (PrimElt a, NumElt a) => Shaped sh a -> Shaped sh a -> a
sdot :: forall a (sh :: [Nat]).
(PrimElt a, NumElt a) =>
Shaped sh a -> Shaped sh a -> a
sdot = (Mixed (MapJust sh) a -> Mixed (MapJust sh) a -> a)
-> Shaped sh a -> Shaped sh a -> a
forall a b. Coercible a b => a -> b
coerce Mixed (MapJust sh) a -> Mixed (MapJust sh) a -> a
forall a (sh :: [Maybe Nat]).
(PrimElt a, NumElt a) =>
Mixed sh a -> Mixed sh a -> a
mdot

stoXArrayPrimP :: Shaped sh (Primitive a) -> (ShS sh, XArray (MapJust sh) a)
stoXArrayPrimP :: forall (sh :: [Nat]) a.
Shaped sh (Primitive a) -> (ShS sh, XArray (MapJust sh) a)
stoXArrayPrimP (Shaped Mixed (MapJust sh) (Primitive a)
arr) = (IShX (MapJust sh) -> ShS sh)
-> (IShX (MapJust sh), XArray (MapJust sh) a)
-> (ShS sh, XArray (MapJust sh) a)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first IShX (MapJust sh) -> ShS sh
forall (sh :: [Nat]) i. ShX (MapJust sh) i -> ShS sh
shsFromShX (Mixed (MapJust sh) (Primitive a)
-> (IShX (MapJust sh), XArray (MapJust sh) a)
forall (sh :: [Maybe Nat]) a.
Mixed sh (Primitive a) -> (IShX sh, XArray sh a)
mtoXArrayPrimP Mixed (MapJust sh) (Primitive a)
arr)

stoXArrayPrim :: PrimElt a => Shaped sh a -> (ShS sh, XArray (MapJust sh) a)
stoXArrayPrim :: forall a (sh :: [Nat]).
PrimElt a =>
Shaped sh a -> (ShS sh, XArray (MapJust sh) a)
stoXArrayPrim (Shaped Mixed (MapJust sh) a
arr) = (IShX (MapJust sh) -> ShS sh)
-> (IShX (MapJust sh), XArray (MapJust sh) a)
-> (ShS sh, XArray (MapJust sh) a)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first IShX (MapJust sh) -> ShS sh
forall (sh :: [Nat]) i. ShX (MapJust sh) i -> ShS sh
shsFromShX (Mixed (MapJust sh) a -> (IShX (MapJust sh), XArray (MapJust sh) a)
forall a (sh :: [Maybe Nat]).
PrimElt a =>
Mixed sh a -> (IShX sh, XArray sh a)
mtoXArrayPrim Mixed (MapJust sh) a
arr)

sfromXArrayPrimP :: ShS sh -> XArray (MapJust sh) a -> Shaped sh (Primitive a)
sfromXArrayPrimP :: forall (sh :: [Nat]) a.
ShS sh -> XArray (MapJust sh) a -> Shaped sh (Primitive a)
sfromXArrayPrimP ShS sh
sh XArray (MapJust sh) a
arr = Mixed (MapJust sh) (Primitive a) -> Shaped sh (Primitive a)
forall (sh :: [Nat]) a. Mixed (MapJust sh) a -> Shaped sh a
Shaped (StaticShX (MapJust sh)
-> XArray (MapJust sh) a -> Mixed (MapJust sh) (Primitive a)
forall (sh :: [Maybe Nat]) a.
StaticShX sh -> XArray sh a -> Mixed sh (Primitive a)
mfromXArrayPrimP (ShX (MapJust sh) Int -> StaticShX (MapJust sh)
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX (ShS sh -> ShX (MapJust sh) Int
forall (sh :: [Nat]). ShS sh -> IShX (MapJust sh)
shxFromShS ShS sh
sh)) XArray (MapJust sh) a
arr)

sfromXArrayPrim :: PrimElt a => ShS sh -> XArray (MapJust sh) a -> Shaped sh a
sfromXArrayPrim :: forall a (sh :: [Nat]).
PrimElt a =>
ShS sh -> XArray (MapJust sh) a -> Shaped sh a
sfromXArrayPrim ShS sh
sh XArray (MapJust sh) a
arr = Mixed (MapJust sh) a -> Shaped sh a
forall (sh :: [Nat]) a. Mixed (MapJust sh) a -> Shaped sh a
Shaped (StaticShX (MapJust sh)
-> XArray (MapJust sh) a -> Mixed (MapJust sh) a
forall a (sh :: [Maybe Nat]).
PrimElt a =>
StaticShX sh -> XArray sh a -> Mixed sh a
mfromXArrayPrim (ShX (MapJust sh) Int -> StaticShX (MapJust sh)
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX (ShS sh -> ShX (MapJust sh) Int
forall (sh :: [Nat]). ShS sh -> IShX (MapJust sh)
shxFromShS ShS sh
sh)) XArray (MapJust sh) a
arr)

sfromPrimitive :: PrimElt a => Shaped sh (Primitive a) -> Shaped sh a
sfromPrimitive :: forall a (sh :: [Nat]).
PrimElt a =>
Shaped sh (Primitive a) -> Shaped sh a
sfromPrimitive (Shaped Mixed (MapJust sh) (Primitive a)
arr) = Mixed (MapJust sh) a -> Shaped sh a
forall (sh :: [Nat]) a. Mixed (MapJust sh) a -> Shaped sh a
Shaped (Mixed (MapJust sh) (Primitive a) -> Mixed (MapJust sh) a
forall (sh :: [Maybe Nat]). Mixed sh (Primitive a) -> Mixed sh a
forall a (sh :: [Maybe Nat]).
PrimElt a =>
Mixed sh (Primitive a) -> Mixed sh a
fromPrimitive Mixed (MapJust sh) (Primitive a)
arr)

stoPrimitive :: PrimElt a => Shaped sh a -> Shaped sh (Primitive a)
stoPrimitive :: forall a (sh :: [Nat]).
PrimElt a =>
Shaped sh a -> Shaped sh (Primitive a)
stoPrimitive (Shaped Mixed (MapJust sh) a
arr) = Mixed (MapJust sh) (Primitive a) -> Shaped sh (Primitive a)
forall (sh :: [Nat]) a. Mixed (MapJust sh) a -> Shaped sh a
Shaped (Mixed (MapJust sh) a -> Mixed (MapJust sh) (Primitive a)
forall (sh :: [Maybe Nat]). Mixed sh a -> Mixed sh (Primitive a)
forall a (sh :: [Maybe Nat]).
PrimElt a =>
Mixed sh a -> Mixed sh (Primitive a)
toPrimitive Mixed (MapJust sh) a
arr)