{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeAbstractions #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeOperators #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
module Data.Array.Nested.Convert (
  -- * Shape\/index\/list casting functions
  -- ** To ranked
  ixrFromIxS, ixrFromIxX, shrFromShS, shrFromShX, shrFromShX2,
  listrCast, ixrCast, shrCast,
  -- ** To shaped
  ixsFromIxR, ixsFromIxR', ixsFromIxX, ixsFromIxX', withShsFromShR, shsFromShX, withShsFromShX, shsFromSSX,
  ixsCast,
  -- ** To mixed
  ixxFromIxR, ixxFromIxS, shxFromShR, shxFromShS,
  ixxCast, shxCast, shxCast',

  -- * Array conversions
  convert,
  Conversion(..),

  -- * Special cases of array conversions
  --
  -- | These functions can all be implemented using 'convert' in some way,
  -- but some have fewer constraints.
  rtoMixed, rcastToMixed, rcastToShaped,
  stoMixed, scastToMixed, stoRanked,
  mcast, mcastToShaped, mtoRanked,
) where

import Control.Category
import Data.Proxy
import Data.Type.Equality
import GHC.TypeLits

import Data.Array.Nested.Lemmas
import Data.Array.Nested.Mixed
import Data.Array.Nested.Mixed.Shape
import Data.Array.Nested.Ranked.Base
import Data.Array.Nested.Ranked.Shape
import Data.Array.Nested.Shaped.Base
import Data.Array.Nested.Shaped.Shape
import Data.Array.Nested.Types

-- * Shape or index or list casting functions

-- * To ranked

ixrFromIxS :: IxS sh i -> IxR (Rank sh) i
ixrFromIxS :: forall (sh :: [Nat]) i. IxS sh i -> IxR (Rank sh) i
ixrFromIxS IxS sh i
ZIS = IxR 0 i
IxR (Rank sh) i
forall (n :: Nat) i. (n ~ 0) => IxR n i
ZIR
ixrFromIxS (i
i :.$ IxS sh i
ix) = i
i i -> IxR (Rank sh) i -> IxR (Rank sh + 1) i
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1) ~ n1) =>
i -> IxR n i -> IxR n1 i
:.: IxS sh i -> IxR (Rank sh) i
forall (sh :: [Nat]) i. IxS sh i -> IxR (Rank sh) i
ixrFromIxS IxS sh i
ix

ixrFromIxX :: IxX sh i -> IxR (Rank sh) i
ixrFromIxX :: forall (sh :: [Maybe Nat]) i. IxX sh i -> IxR (Rank sh) i
ixrFromIxX IxX sh i
ZIX = IxR 0 i
IxR (Rank sh) i
forall (n :: Nat) i. (n ~ 0) => IxR n i
ZIR
ixrFromIxX (i
n :.% IxX sh i
idx) = i
n i -> IxR (Rank sh) i -> IxR (Rank sh + 1) i
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1) ~ n1) =>
i -> IxR n i -> IxR n1 i
:.: IxX sh i -> IxR (Rank sh) i
forall (sh :: [Maybe Nat]) i. IxX sh i -> IxR (Rank sh) i
ixrFromIxX IxX sh i
idx

shrFromShS :: ShS sh -> IShR (Rank sh)
shrFromShS :: forall (sh :: [Nat]). ShS sh -> IShR (Rank sh)
shrFromShS ShS sh
ZSS = ShR 0 Int
ShR (Rank sh) Int
forall (n :: Nat) i. (n ~ 0) => ShR n i
ZSR
shrFromShS (SNat n
n :$$ ShS sh
sh) = SNat n -> Int
forall (n :: Nat). SNat n -> Int
fromSNat' SNat n
n Int -> ShR (Rank sh) Int -> ShR (Rank sh + 1) Int
forall {n1 :: Nat} {i} (n :: Nat).
((n + 1) ~ n1) =>
i -> ShR n i -> ShR n1 i
:$: ShS sh -> ShR (Rank sh) Int
forall (sh :: [Nat]). ShS sh -> IShR (Rank sh)
shrFromShS ShS sh
sh

-- shrFromShX re-exported
-- shrFromShX2 re-exported
-- listrCast re-exported
-- ixrCast re-exported
-- shrCast re-exported

-- * To shaped

-- TODO: these take a ShS because there are KnownNats inside IxS.

ixsFromIxR :: ShS sh -> IxR (Rank sh) i -> IxS sh i
ixsFromIxR :: forall (sh :: [Nat]) i. ShS sh -> IxR (Rank sh) i -> IxS sh i
ixsFromIxR ShS sh
ZSS IxR 0 i
IxR (Rank sh) i
ZIR = IxS sh i
forall (sh :: [Nat]) i. (sh ~ '[]) => IxS sh i
ZIS
ixsFromIxR (SNat n
_ :$$ ShS sh
sh) (i
n :.: IxR n i
idx) = i
n i -> IxS sh i -> IxS sh i
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, (n : sh) ~ sh1) =>
i -> IxS sh i -> IxS sh1 i
:.$ ShS sh -> IxR (Rank sh) i -> IxS sh i
forall (sh :: [Nat]) i. ShS sh -> IxR (Rank sh) i -> IxS sh i
ixsFromIxR ShS sh
sh IxR n i
IxR (Rank sh) i
idx
ixsFromIxR ShS sh
_ IxR (Rank sh) i
_ = [Char] -> IxS sh i
forall a. HasCallStack => [Char] -> a
error [Char]
"unreachable"

-- | Performs a runtime check that @n@ matches @Rank sh@. Equivalent to the
-- following, but more efficient:
--
-- > ixsFromIxR' sh idx = ixsFromIxR sh (ixrCast (shsRank sh) idx)
ixsFromIxR' :: ShS sh -> IxR n i -> IxS sh i
ixsFromIxR' :: forall (sh :: [Nat]) (n :: Nat) i. ShS sh -> IxR n i -> IxS sh i
ixsFromIxR' ShS sh
ZSS IxR n i
ZIR = IxS sh i
forall (sh :: [Nat]) i. (sh ~ '[]) => IxS sh i
ZIS
ixsFromIxR' (SNat n
_ :$$ ShS sh
sh) (i
n :.: IxR n i
idx) = i
n i -> IxS sh i -> IxS sh i
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, (n : sh) ~ sh1) =>
i -> IxS sh i -> IxS sh1 i
:.$ ShS sh -> IxR n i -> IxS sh i
forall (sh :: [Nat]) (n :: Nat) i. ShS sh -> IxR n i -> IxS sh i
ixsFromIxR' ShS sh
sh IxR n i
idx
ixsFromIxR' ShS sh
_ IxR n i
_ = [Char] -> IxS sh i
forall a. HasCallStack => [Char] -> a
error [Char]
"ixsFromIxR': index rank does not match shape rank"

-- TODO: this takes a ShS because there are KnownNats inside IxS.
ixsFromIxX :: ShS sh -> IxX (MapJust sh) i -> IxS sh i
ixsFromIxX :: forall (sh :: [Nat]) i. ShS sh -> IxX (MapJust sh) i -> IxS sh i
ixsFromIxX ShS sh
ZSS IxX '[] i
IxX (MapJust sh) i
ZIX = IxS sh i
forall (sh :: [Nat]) i. (sh ~ '[]) => IxS sh i
ZIS
ixsFromIxX (SNat n
_ :$$ ShS sh
sh) (i
n :.% IxX sh i
idx) = i
n i -> IxS sh i -> IxS sh i
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, (n : sh) ~ sh1) =>
i -> IxS sh i -> IxS sh1 i
:.$ ShS sh -> IxX (MapJust sh) i -> IxS sh i
forall (sh :: [Nat]) i. ShS sh -> IxX (MapJust sh) i -> IxS sh i
ixsFromIxX ShS sh
sh IxX sh i
IxX (MapJust sh) i
idx

-- | Performs a runtime check that @Rank sh'@ match @Rank sh@. Equivalent to
-- the following, but more efficient:
--
-- > ixsFromIxX' sh idx = ixsFromIxX sh (ixxCast (shxFromShS sh) idx)
ixsFromIxX' :: ShS sh -> IxX sh' i -> IxS sh i
ixsFromIxX' :: forall (sh :: [Nat]) (sh' :: [Maybe Nat]) i.
ShS sh -> IxX sh' i -> IxS sh i
ixsFromIxX' ShS sh
ZSS IxX sh' i
ZIX = IxS sh i
forall (sh :: [Nat]) i. (sh ~ '[]) => IxS sh i
ZIS
ixsFromIxX' (SNat n
_ :$$ ShS sh
sh) (i
n :.% IxX sh i
idx) = i
n i -> IxS sh i -> IxS sh i
forall {sh1 :: [Nat]} {i} (n :: Nat) (sh :: [Nat]).
(KnownNat n, (n : sh) ~ sh1) =>
i -> IxS sh i -> IxS sh1 i
:.$ ShS sh -> IxX sh i -> IxS sh i
forall (sh :: [Nat]) (sh' :: [Maybe Nat]) i.
ShS sh -> IxX sh' i -> IxS sh i
ixsFromIxX' ShS sh
sh IxX sh i
idx
ixsFromIxX' ShS sh
_ IxX sh' i
_ = [Char] -> IxS sh i
forall a. HasCallStack => [Char] -> a
error [Char]
"ixsFromIxX': index rank does not match shape rank"

-- | Produce an existential 'ShS' from an 'IShR'.
withShsFromShR :: IShR n -> (forall sh. Rank sh ~ n => ShS sh -> r) -> r
withShsFromShR :: forall (n :: Nat) r.
IShR n -> (forall (sh :: [Nat]). (Rank sh ~ n) => ShS sh -> r) -> r
withShsFromShR ShR n Int
ZSR forall (sh :: [Nat]). (Rank sh ~ n) => ShS sh -> r
k = ShS '[] -> r
forall (sh :: [Nat]). (Rank sh ~ n) => ShS sh -> r
k ShS '[]
forall (sh :: [Nat]). (sh ~ '[]) => ShS sh
ZSS
withShsFromShR (Int
n :$: ShR n Int
sh) forall (sh :: [Nat]). (Rank sh ~ n) => ShS sh -> r
k =
  ShR n Int
-> (forall {sh :: [Nat]}. (Rank sh ~ n) => ShS sh -> r) -> r
forall (n :: Nat) r.
IShR n -> (forall (sh :: [Nat]). (Rank sh ~ n) => ShS sh -> r) -> r
withShsFromShR ShR n Int
sh ((forall {sh :: [Nat]}. (Rank sh ~ n) => ShS sh -> r) -> r)
-> (forall {sh :: [Nat]}. (Rank sh ~ n) => ShS sh -> r) -> r
forall a b. (a -> b) -> a -> b
$ \ShS sh
sh' ->
    Integer -> (forall {n :: Nat}. Maybe (SNat n) -> r) -> r
forall r. Integer -> (forall (n :: Nat). Maybe (SNat n) -> r) -> r
withSomeSNat (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Int @Integer Int
n) ((forall {n :: Nat}. Maybe (SNat n) -> r) -> r)
-> (forall {n :: Nat}. Maybe (SNat n) -> r) -> r
forall a b. (a -> b) -> a -> b
$ \case
      Just sn :: SNat n
sn@SNat n
SNat -> ShS (n : sh) -> r
forall (sh :: [Nat]). (Rank sh ~ n) => ShS sh -> r
k (SNat n
sn 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')
      Maybe (SNat n)
Nothing -> [Char] -> r
forall a. HasCallStack => [Char] -> a
error ([Char] -> r) -> [Char] -> r
forall a b. (a -> b) -> a -> b
$ [Char]
"withShsFromShR: negative dimension size (" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
n [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
")"

-- shsFromShX re-exported

-- | Produce an existential 'ShS' from an 'IShX'. If you already know that
-- @sh'@ is @MapJust@ of something, use 'shsFromShX' instead.
withShsFromShX :: IShX sh' -> (forall sh. Rank sh ~ Rank sh' => ShS sh -> r) -> r
withShsFromShX :: forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]). (Rank sh ~ Rank sh') => ShS sh -> r) -> r
withShsFromShX ShX sh' Int
ZSX forall (sh :: [Nat]). (Rank sh ~ Rank sh') => ShS sh -> r
k = ShS '[] -> r
forall (sh :: [Nat]). (Rank sh ~ Rank sh') => ShS sh -> r
k ShS '[]
forall (sh :: [Nat]). (sh ~ '[]) => ShS sh
ZSS
withShsFromShX (SKnown sn :: SNat n1
sn@SNat n1
SNat :$% ShX sh Int
sh) forall (sh :: [Nat]). (Rank sh ~ Rank sh') => ShS sh -> r
k =
  ShX sh Int
-> (forall {sh :: [Nat]}. (Rank sh ~ Rank sh) => ShS sh -> r) -> r
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]). (Rank sh ~ Rank sh') => ShS sh -> r) -> r
withShsFromShX ShX sh Int
sh ((forall {sh :: [Nat]}. (Rank sh ~ Rank sh) => ShS sh -> r) -> r)
-> (forall {sh :: [Nat]}. (Rank sh ~ Rank sh) => ShS sh -> r) -> r
forall a b. (a -> b) -> a -> b
$ \ShS sh
sh' ->
    ShS (n1 : sh) -> r
forall (sh :: [Nat]). (Rank sh ~ Rank sh') => ShS sh -> r
k (SNat n1
sn SNat n1 -> ShS sh -> ShS (n1 : sh)
forall {sh1 :: [Nat]} (n :: Nat) (sh :: [Nat]).
(KnownNat n, (n : sh) ~ sh1) =>
SNat n -> ShS sh -> ShS sh1
:$$ ShS sh
sh')
withShsFromShX (SUnknown Int
n :$% ShX sh Int
sh) forall (sh :: [Nat]). (Rank sh ~ Rank sh') => ShS sh -> r
k =
  ShX sh Int
-> (forall {sh :: [Nat]}. (Rank sh ~ Rank sh) => ShS sh -> r) -> r
forall (sh' :: [Maybe Nat]) r.
IShX sh'
-> (forall (sh :: [Nat]). (Rank sh ~ Rank sh') => ShS sh -> r) -> r
withShsFromShX ShX sh Int
sh ((forall {sh :: [Nat]}. (Rank sh ~ Rank sh) => ShS sh -> r) -> r)
-> (forall {sh :: [Nat]}. (Rank sh ~ Rank sh) => ShS sh -> r) -> r
forall a b. (a -> b) -> a -> b
$ \ShS sh
sh' ->
    Integer -> (forall {n :: Nat}. Maybe (SNat n) -> r) -> r
forall r. Integer -> (forall (n :: Nat). Maybe (SNat n) -> r) -> r
withSomeSNat (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Int @Integer Int
n) ((forall {n :: Nat}. Maybe (SNat n) -> r) -> r)
-> (forall {n :: Nat}. Maybe (SNat n) -> r) -> r
forall a b. (a -> b) -> a -> b
$ \case
      Just sn :: SNat n
sn@SNat n
SNat -> ShS (n : sh) -> r
forall (sh :: [Nat]). (Rank sh ~ Rank sh') => ShS sh -> r
k (SNat n
sn 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')
      Maybe (SNat n)
Nothing -> [Char] -> r
forall a. HasCallStack => [Char] -> a
error ([Char] -> r) -> [Char] -> r
forall a b. (a -> b) -> a -> b
$ [Char]
"withShsFromShX: negative SUnknown dimension size (" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
n [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
")"

shsFromSSX :: StaticShX (MapJust sh) -> ShS sh
shsFromSSX :: forall (sh :: [Nat]). StaticShX (MapJust sh) -> ShS sh
shsFromSSX = ShX (MapJust sh) Any -> ShS sh
forall (sh :: [Nat]) i. ShX (MapJust sh) i -> ShS sh
shsFromShX (ShX (MapJust sh) Any -> ShS sh)
-> (StaticShX (MapJust sh) -> ShX (MapJust sh) Any)
-> StaticShX (MapJust sh)
-> ShS sh
forall b c a. (b -> c) -> (a -> b) -> a -> c
Prelude.. StaticShX (MapJust sh) -> ShX (MapJust sh) Any
forall (sh :: [Nat]) i.
StaticShX (MapJust sh) -> ShX (MapJust sh) i
shxFromSSX

-- ixsCast re-exported

-- * To mixed

ixxFromIxR :: IxR n i -> IxX (Replicate n Nothing) i
ixxFromIxR :: forall (n :: Nat) i. IxR n i -> IxX (Replicate n 'Nothing) i
ixxFromIxR IxR n i
ZIR = IxX '[] i
IxX (Replicate n 'Nothing) i
forall (sh :: [Maybe Nat]) i. (sh ~ '[]) => IxX sh i
ZIX
ixxFromIxR (i
n :.: (IxR n i
idx :: IxR m i)) =
  (IxX ('Nothing : Replicate n 'Nothing) i
 :~: IxX (Replicate n 'Nothing) i)
-> IxX ('Nothing : Replicate n 'Nothing) i
-> IxX (Replicate n 'Nothing) i
forall a b. (a :~: b) -> a -> b
castWith (forall {k1} {k2} {k3} (f :: k1 -> k2 -> k3) (c :: k2) (a :: k1)
       (b :: k1).
(a :~: b) -> f a c :~: f b c
forall (f :: [Maybe Nat] -> * -> *) c (a :: [Maybe Nat])
       (b :: [Maybe Nat]).
(a :~: b) -> f a c :~: f b c
subst2 @IxX @i (forall (a2 :: Maybe Nat) (n :: Nat).
(a2 : Replicate n a2) :~: Replicate (n + 1) a2
forall {a1} (a2 :: a1) (n :: Nat).
(a2 : Replicate n a2) :~: Replicate (n + 1) a2
lemReplicateSucc @(Nothing @Nat) @m))
    (i
n i
-> IxX (Replicate n 'Nothing) i
-> IxX ('Nothing : Replicate n 'Nothing) i
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
i -> IxX sh i -> IxX sh1 i
:.% IxR n i -> IxX (Replicate n 'Nothing) i
forall (n :: Nat) i. IxR n i -> IxX (Replicate n 'Nothing) i
ixxFromIxR IxR n i
idx)

ixxFromIxS :: IxS sh i -> IxX (MapJust sh) i
ixxFromIxS :: forall (sh :: [Nat]) i. IxS sh i -> IxX (MapJust sh) i
ixxFromIxS IxS sh i
ZIS = IxX '[] i
IxX (MapJust sh) i
forall (sh :: [Maybe Nat]) i. (sh ~ '[]) => IxX sh i
ZIX
ixxFromIxS (i
n :.$ IxS sh i
sh) = i
n i -> IxX (MapJust sh) i -> IxX ('Just n : MapJust sh) i
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
i -> IxX sh i -> IxX sh1 i
:.% IxS sh i -> IxX (MapJust sh) i
forall (sh :: [Nat]) i. IxS sh i -> IxX (MapJust sh) i
ixxFromIxS IxS sh i
sh

shxFromShR :: ShR n i -> ShX (Replicate n Nothing) i
shxFromShR :: forall (n :: Nat) i. ShR n i -> ShX (Replicate n 'Nothing) i
shxFromShR ShR n i
ZSR = ShX '[] i
ShX (Replicate n 'Nothing) i
forall (sh :: [Maybe Nat]) i. (sh ~ '[]) => ShX sh i
ZSX
shxFromShR (i
n :$: (ShR n i
idx :: ShR m i)) =
  (ShX ('Nothing : Replicate n 'Nothing) i
 :~: ShX (Replicate n 'Nothing) i)
-> ShX ('Nothing : Replicate n 'Nothing) i
-> ShX (Replicate n 'Nothing) i
forall a b. (a :~: b) -> a -> b
castWith (forall {k1} {k2} {k3} (f :: k1 -> k2 -> k3) (c :: k2) (a :: k1)
       (b :: k1).
(a :~: b) -> f a c :~: f b c
forall (f :: [Maybe Nat] -> * -> *) c (a :: [Maybe Nat])
       (b :: [Maybe Nat]).
(a :~: b) -> f a c :~: f b c
subst2 @ShX @i (forall (a2 :: Maybe Nat) (n :: Nat).
(a2 : Replicate n a2) :~: Replicate (n + 1) a2
forall {a1} (a2 :: a1) (n :: Nat).
(a2 : Replicate n a2) :~: Replicate (n + 1) a2
lemReplicateSucc @(Nothing @Nat) @m))
    (i -> SMayNat i SNat 'Nothing
forall {k} i (f :: k -> *). i -> SMayNat i f 'Nothing
SUnknown i
n SMayNat i SNat 'Nothing
-> ShX (Replicate n 'Nothing) i
-> ShX ('Nothing : Replicate n 'Nothing) i
forall {sh1 :: [Maybe Nat]} {i} (n :: Maybe Nat)
       (sh :: [Maybe Nat]).
((n : sh) ~ sh1) =>
SMayNat i SNat n -> ShX sh i -> ShX sh1 i
:$% ShR n i -> ShX (Replicate n 'Nothing) i
forall (n :: Nat) i. ShR n i -> ShX (Replicate n 'Nothing) i
shxFromShR ShR n i
idx)

shxFromShS :: ShS sh -> IShX (MapJust sh)
shxFromShS :: forall (sh :: [Nat]). ShS sh -> IShX (MapJust sh)
shxFromShS ShS sh
ZSS = ShX '[] Int
ShX (MapJust sh) Int
forall (sh :: [Maybe Nat]) i. (sh ~ '[]) => ShX sh i
ZSX
shxFromShS (SNat n
n :$$ ShS sh
sh) = SNat n -> SMayNat Int SNat ('Just n)
forall {k} (f :: k -> *) (n1 :: k) i.
f n1 -> SMayNat i f ('Just n1)
SKnown SNat n
n SMayNat Int SNat ('Just n)
-> ShX (MapJust sh) Int -> ShX ('Just n : MapJust 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
:$% ShS sh -> ShX (MapJust sh) Int
forall (sh :: [Nat]). ShS sh -> IShX (MapJust sh)
shxFromShS ShS sh
sh

-- ixxCast re-exported
-- shxCast re-exported
-- shxCast' re-exported


-- * Array conversions

-- | The constructors that perform runtime shape checking are marked with a
-- tick (@'@): 'ConvXS'' and 'ConvXX''. For the other constructors, the types
-- ensure that the shapes are already compatible. To convert between 'Ranked'
-- and 'Shaped', go via 'Mixed'.
--
-- The guiding principle behind 'Conversion' is that it should represent the
-- array restructurings, or perhaps re-presentations, that do not change the
-- underlying 'XArray's. This leads to the inclusion of some operations that do
-- not look like simple conversions (casts) at first glance, like 'ConvZip'.
--
-- /Note/: Haddock gleefully renames type variables in constructors so that
-- they match the data type head as much as possible. See the source for a more
-- readable presentation of this data type.
data Conversion a b where
  ConvId  :: Conversion a a
  ConvCmp :: Conversion b c -> Conversion a b -> Conversion a c

  ConvRX  :: Conversion (Ranked n a) (Mixed (Replicate n Nothing) a)
  ConvSX  :: Conversion (Shaped sh a) (Mixed (MapJust sh) a)

  ConvXR  :: Elt a
          => Conversion (Mixed sh a) (Ranked (Rank sh) a)
  ConvXS  :: Conversion (Mixed (MapJust sh) a) (Shaped sh a)
  ConvXS' :: (Rank sh ~ Rank sh', Elt a)
          => ShS sh'
          -> Conversion (Mixed sh a) (Shaped sh' a)

  ConvXX' :: (Rank sh ~ Rank sh', Elt a)
          => StaticShX sh'
          -> Conversion (Mixed sh a) (Mixed sh' a)

  ConvRR  :: Conversion a b
          -> Conversion (Ranked n a) (Ranked n b)
  ConvSS  :: Conversion a b
          -> Conversion (Shaped sh a) (Shaped sh b)
  ConvXX  :: Conversion a b
          -> Conversion (Mixed sh a) (Mixed sh b)
  ConvT2  :: Conversion a a'
          -> Conversion b b'
          -> Conversion (a, b) (a', b')

  Conv0X  :: Elt a
          => Conversion a (Mixed '[] a)
  ConvX0  :: Conversion (Mixed '[] a) a

  ConvNest   :: Elt a => StaticShX sh
             -> Conversion (Mixed (sh ++ sh') a) (Mixed sh (Mixed sh' a))
  ConvUnnest :: Conversion (Mixed sh (Mixed sh' a)) (Mixed (sh ++ sh') a)

  ConvZip   :: (Elt a, Elt b)
            => Conversion (Mixed sh a, Mixed sh b) (Mixed sh (a, b))
  ConvUnzip :: (Elt a, Elt b)
            => Conversion (Mixed sh (a, b)) (Mixed sh a, Mixed sh b)
deriving instance Show (Conversion a b)

instance Category Conversion where
  id :: forall a. Conversion a a
id = Conversion a a
forall a. Conversion a a
ConvId
  . :: forall b c a. Conversion b c -> Conversion a b -> Conversion a c
(.) = Conversion b c -> Conversion a b -> Conversion a c
forall b c a. Conversion b c -> Conversion a b -> Conversion a c
ConvCmp

convert :: (Elt a, Elt b) => Conversion a b -> a -> b
convert :: forall a b. (Elt a, Elt b) => Conversion a b -> a -> b
convert = \Conversion a b
c a
x -> Mixed '[] b -> b
forall a. Elt a => Mixed '[] a -> a
munScalar (Conversion a b -> Mixed '[] a -> Mixed '[] b
forall a b (esh :: [Maybe Nat]).
Conversion a b -> Mixed esh a -> Mixed esh b
go Conversion a b
c (a -> Mixed '[] a
forall a. Elt a => a -> Mixed '[] a
mscalar a
x))
  where
    -- The 'esh' is the extension shape: the conversion happens under a whole
    -- bunch of additional dimensions that it does not touch. These dimensions
    -- are 'esh'.
    -- The strategy is to unwind step-by-step to a large Mixed array, and to
    -- perform the required checks and conversions when re-nesting back up.
    go :: Conversion a b -> Mixed esh a -> Mixed esh b
    go :: forall a b (esh :: [Maybe Nat]).
Conversion a b -> Mixed esh a -> Mixed esh b
go Conversion a b
ConvId Mixed esh a
x = Mixed esh a
Mixed esh b
x
    go (ConvCmp Conversion b b
c1 Conversion a b
c2) Mixed esh a
x = Conversion b b -> Mixed esh b -> Mixed esh b
forall a b (esh :: [Maybe Nat]).
Conversion a b -> Mixed esh a -> Mixed esh b
go Conversion b b
c1 (Conversion a b -> Mixed esh a -> Mixed esh b
forall a b (esh :: [Maybe Nat]).
Conversion a b -> Mixed esh a -> Mixed esh b
go Conversion a b
c2 Mixed esh a
x)
    go Conversion a b
ConvRX (M_Ranked Mixed esh (Mixed (Replicate n 'Nothing) a)
x) = Mixed esh b
Mixed esh (Mixed (Replicate n 'Nothing) a)
x
    go Conversion a b
ConvSX (M_Shaped Mixed esh (Mixed (MapJust sh) a)
x) = Mixed esh b
Mixed esh (Mixed (MapJust sh) a)
x
    go (ConvXR @_ @sh) (M_Nest @esh IShX esh
esh Mixed (esh ++ sh) a
x)
      | Rank (esh ++ sh) :~: Rank (esh ++ Replicate (Rank sh) 'Nothing)
Refl <- Proxy esh
-> Proxy sh
-> Rank (esh ++ sh) :~: Rank (esh ++ Replicate (Rank sh) 'Nothing)
forall {a} (esh :: [Maybe a]) (sh :: [Maybe a]).
Proxy esh
-> Proxy sh
-> Rank (esh ++ sh) :~: Rank (esh ++ Replicate (Rank sh) 'Nothing)
lemRankAppRankEqRepNo (forall (t :: [Maybe Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @esh) (forall (t :: [Maybe Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @sh)
      = let ssx' :: StaticShX (esh ++ Replicate (Rank sh) 'Nothing)
ssx' = StaticShX esh
-> StaticShX (Replicate (Rank sh) 'Nothing)
-> StaticShX (esh ++ Replicate (Rank sh) 'Nothing)
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
StaticShX sh -> StaticShX sh' -> StaticShX (sh ++ sh')
ssxAppend (IShX esh -> StaticShX esh
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX IShX esh
esh)
                             (SNat (Rank sh) -> StaticShX (Replicate (Rank sh) 'Nothing)
forall (n :: Nat). SNat n -> StaticShX (Replicate n 'Nothing)
ssxReplicate (ShX sh Int -> SNat (Rank sh)
forall (sh :: [Maybe Nat]) i. ShX sh i -> SNat (Rank sh)
shxRank (forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) i.
StaticShX sh -> ShX (sh ++ sh') i -> ShX sh' i
shxDropSSX @esh @sh (IShX esh -> StaticShX esh
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX IShX esh
esh) (Mixed (esh ++ sh) a -> ShX (esh ++ sh) Int
forall (sh :: [Maybe Nat]). Mixed sh a -> IShX sh
forall a (sh :: [Maybe Nat]). Elt a => Mixed sh a -> IShX sh
mshape Mixed (esh ++ sh) a
x))))
        in Mixed esh (Mixed (Replicate (Rank sh) 'Nothing) a)
-> Mixed esh (Ranked (Rank sh) a)
forall (sh :: [Maybe Nat]) (n :: Nat) a.
Mixed sh (Mixed (Replicate n 'Nothing) a) -> Mixed sh (Ranked n a)
M_Ranked (IShX esh
-> Mixed (esh ++ Replicate (Rank sh) 'Nothing) a
-> Mixed esh (Mixed (Replicate (Rank sh) 'Nothing) a)
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
IShX sh1 -> Mixed (sh1 ++ sh2) a -> Mixed sh1 (Mixed sh2 a)
M_Nest IShX esh
esh (StaticShX (esh ++ Replicate (Rank sh) 'Nothing)
-> Mixed (esh ++ sh) a
-> Mixed (esh ++ Replicate (Rank sh) 'Nothing) a
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
(Rank sh1 ~ Rank sh2, Elt a) =>
StaticShX sh2 -> Mixed sh1 a -> Mixed sh2 a
mcast StaticShX (esh ++ Replicate (Rank sh) 'Nothing)
ssx' Mixed (esh ++ sh) a
x))
    go Conversion a b
ConvXS (M_Nest IShX esh
esh Mixed (esh ++ MapJust sh) a
x) = Mixed esh (Mixed (MapJust sh) a) -> Mixed esh (Shaped sh a)
forall (sh :: [Maybe Nat]) (sh' :: [Nat]) a.
Mixed sh (Mixed (MapJust sh') a) -> Mixed sh (Shaped sh' a)
M_Shaped (IShX esh
-> Mixed (esh ++ MapJust sh) a -> Mixed esh (Mixed (MapJust sh) a)
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
IShX sh1 -> Mixed (sh1 ++ sh2) a -> Mixed sh1 (Mixed sh2 a)
M_Nest IShX esh
esh Mixed (esh ++ MapJust sh) a
x)
    go (ConvXS' @sh @sh' ShS sh'
sh') (M_Nest @esh IShX esh
esh Mixed (esh ++ sh) a
x)
      | Rank (esh ++ sh) :~: Rank (esh ++ MapJust sh')
Refl <- Proxy esh
-> Proxy sh
-> Proxy sh'
-> Rank (esh ++ sh) :~: Rank (esh ++ MapJust sh')
forall {a} (sh :: [Maybe a]) (sh' :: [a]) (esh :: [Maybe a]).
(Rank sh ~ Rank sh') =>
Proxy esh
-> Proxy sh
-> Proxy sh'
-> Rank (esh ++ sh) :~: Rank (esh ++ MapJust sh')
lemRankAppRankEqMapJust (forall (t :: [Maybe Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @esh) (forall (t :: [Maybe Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @sh) (forall (t :: [Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @sh')
      = Mixed esh (Mixed (MapJust sh') a) -> Mixed esh (Shaped sh' a)
forall (sh :: [Maybe Nat]) (sh' :: [Nat]) a.
Mixed sh (Mixed (MapJust sh') a) -> Mixed sh (Shaped sh' a)
M_Shaped (IShX esh
-> Mixed (esh ++ MapJust sh') a
-> Mixed esh (Mixed (MapJust sh') a)
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
IShX sh1 -> Mixed (sh1 ++ sh2) a -> Mixed sh1 (Mixed sh2 a)
M_Nest IShX esh
esh (StaticShX (esh ++ MapJust sh')
-> Mixed (esh ++ sh) a -> Mixed (esh ++ MapJust sh') a
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
(Rank sh1 ~ Rank sh2, Elt a) =>
StaticShX sh2 -> Mixed sh1 a -> Mixed sh2 a
mcast (ShX (esh ++ MapJust sh') Int -> StaticShX (esh ++ MapJust sh')
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX (IShX esh -> ShX (MapJust sh') Int -> ShX (esh ++ MapJust sh') Int
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]) i.
ShX sh i -> ShX sh' i -> ShX (sh ++ sh') i
shxAppend IShX esh
esh (ShS sh' -> ShX (MapJust sh') Int
forall (sh :: [Nat]). ShS sh -> IShX (MapJust sh)
shxFromShS ShS sh'
sh')))
                                    Mixed (esh ++ sh) a
x))
    go (ConvXX' @sh @sh' StaticShX sh'
ssx) (M_Nest @esh IShX esh
esh Mixed (esh ++ sh) a
x)
      | Rank (esh ++ sh) :~: Rank (esh ++ sh')
Refl <- Proxy esh
-> Proxy sh -> Proxy sh' -> Rank (esh ++ sh) :~: Rank (esh ++ sh')
forall {a} (sh :: [a]) (sh' :: [a]) (esh :: [a]).
(Rank sh ~ Rank sh') =>
Proxy esh
-> Proxy sh -> Proxy sh' -> Rank (esh ++ sh) :~: Rank (esh ++ sh')
lemRankAppRankEq (forall (t :: [Maybe Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @esh) (forall (t :: [Maybe Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @sh) (forall (t :: [Maybe Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @sh')
      = IShX esh -> Mixed (esh ++ sh') a -> Mixed esh (Mixed sh' a)
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
IShX sh1 -> Mixed (sh1 ++ sh2) a -> Mixed sh1 (Mixed sh2 a)
M_Nest IShX esh
esh (Mixed (esh ++ sh') a -> Mixed esh (Mixed sh' a))
-> Mixed (esh ++ sh') a -> Mixed esh (Mixed sh' a)
forall a b. (a -> b) -> a -> b
$ StaticShX (esh ++ sh')
-> Mixed (esh ++ sh) a -> Mixed (esh ++ sh') a
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
(Rank sh1 ~ Rank sh2, Elt a) =>
StaticShX sh2 -> Mixed sh1 a -> Mixed sh2 a
mcast (IShX esh -> StaticShX esh
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX IShX esh
esh StaticShX esh -> StaticShX sh' -> StaticShX (esh ++ sh')
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
StaticShX sh -> StaticShX sh' -> StaticShX (sh ++ sh')
`ssxAppend` StaticShX sh'
ssx) Mixed (esh ++ sh) a
x
    go (ConvRR Conversion a b
c) (M_Ranked (M_Nest IShX esh
esh Mixed (esh ++ Replicate n 'Nothing) a
x)) = Mixed esh (Mixed (Replicate n 'Nothing) b)
-> Mixed esh (Ranked n b)
forall (sh :: [Maybe Nat]) (n :: Nat) a.
Mixed sh (Mixed (Replicate n 'Nothing) a) -> Mixed sh (Ranked n a)
M_Ranked (IShX esh
-> Mixed (esh ++ Replicate n 'Nothing) b
-> Mixed esh (Mixed (Replicate n 'Nothing) b)
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
IShX sh1 -> Mixed (sh1 ++ sh2) a -> Mixed sh1 (Mixed sh2 a)
M_Nest IShX esh
esh (Conversion a b
-> Mixed (esh ++ Replicate n 'Nothing) a
-> Mixed (esh ++ Replicate n 'Nothing) b
forall a b (esh :: [Maybe Nat]).
Conversion a b -> Mixed esh a -> Mixed esh b
go Conversion a b
c Mixed (esh ++ Replicate n 'Nothing) a
x))
    go (ConvSS Conversion a b
c) (M_Shaped (M_Nest IShX esh
esh Mixed (esh ++ MapJust sh) a
x)) = Mixed esh (Mixed (MapJust sh) b) -> Mixed esh (Shaped sh b)
forall (sh :: [Maybe Nat]) (sh' :: [Nat]) a.
Mixed sh (Mixed (MapJust sh') a) -> Mixed sh (Shaped sh' a)
M_Shaped (IShX esh
-> Mixed (esh ++ MapJust sh) b -> Mixed esh (Mixed (MapJust sh) b)
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
IShX sh1 -> Mixed (sh1 ++ sh2) a -> Mixed sh1 (Mixed sh2 a)
M_Nest IShX esh
esh (Conversion a b
-> Mixed (esh ++ MapJust sh) a -> Mixed (esh ++ MapJust sh) b
forall a b (esh :: [Maybe Nat]).
Conversion a b -> Mixed esh a -> Mixed esh b
go Conversion a b
c Mixed (esh ++ MapJust sh) a
x))
    go (ConvXX Conversion a b
c) (M_Nest IShX esh
esh Mixed (esh ++ sh) a
x) = IShX esh -> Mixed (esh ++ sh) b -> Mixed esh (Mixed sh b)
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
IShX sh1 -> Mixed (sh1 ++ sh2) a -> Mixed sh1 (Mixed sh2 a)
M_Nest IShX esh
esh (Conversion a b -> Mixed (esh ++ sh) a -> Mixed (esh ++ sh) b
forall a b (esh :: [Maybe Nat]).
Conversion a b -> Mixed esh a -> Mixed esh b
go Conversion a b
c Mixed (esh ++ sh) a
x)
    go (ConvT2 Conversion a a'
c1 Conversion b b'
c2) (M_Tup2 Mixed esh a
x1 Mixed esh b
x2) = Mixed esh a' -> Mixed esh b' -> Mixed esh (a', b')
forall (sh :: [Maybe Nat]) a b.
Mixed sh a -> Mixed sh b -> Mixed sh (a, b)
M_Tup2 (Conversion a a' -> Mixed esh a -> Mixed esh a'
forall a b (esh :: [Maybe Nat]).
Conversion a b -> Mixed esh a -> Mixed esh b
go Conversion a a'
c1 Mixed esh a
x1) (Conversion b b' -> Mixed esh b -> Mixed esh b'
forall a b (esh :: [Maybe Nat]).
Conversion a b -> Mixed esh a -> Mixed esh b
go Conversion b b'
c2 Mixed esh b
x2)
    go Conversion a b
Conv0X (Mixed esh a
x :: Mixed esh a)
      | (esh ++ '[]) :~: esh
Refl <- forall (l :: [Maybe Nat]). (l ++ '[]) :~: l
forall {a} (l :: [a]). (l ++ '[]) :~: l
lemAppNil @esh
      = IShX esh -> Mixed (esh ++ '[]) a -> Mixed esh (Mixed '[] a)
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
IShX sh1 -> Mixed (sh1 ++ sh2) a -> Mixed sh1 (Mixed sh2 a)
M_Nest (Mixed esh a -> IShX esh
forall (sh :: [Maybe Nat]). Mixed sh a -> IShX sh
forall a (sh :: [Maybe Nat]). Elt a => Mixed sh a -> IShX sh
mshape Mixed esh a
x) Mixed esh a
Mixed (esh ++ '[]) a
x
    go Conversion a b
ConvX0 (M_Nest @esh IShX esh
_ Mixed (esh ++ '[]) b
x)
      | (esh ++ '[]) :~: esh
Refl <- forall (l :: [Maybe Nat]). (l ++ '[]) :~: l
forall {a} (l :: [a]). (l ++ '[]) :~: l
lemAppNil @esh
      = Mixed esh b
Mixed (esh ++ '[]) b
x
    go (ConvNest @_ @sh @sh' StaticShX sh
ssh) (M_Nest @esh IShX esh
esh Mixed (esh ++ (sh ++ sh')) a
x)
      | ((esh ++ sh) ++ sh') :~: (esh ++ (sh ++ sh'))
Refl <- Proxy esh
-> Proxy sh
-> Proxy sh'
-> ((esh ++ sh) ++ sh') :~: (esh ++ (sh ++ sh'))
forall {a1} (a2 :: [a1]) (b :: [a1]) (c :: [a1]).
Proxy a2
-> Proxy b -> Proxy c -> ((a2 ++ b) ++ c) :~: (a2 ++ (b ++ c))
lemAppAssoc (forall (t :: [Maybe Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @esh) (forall (t :: [Maybe Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @sh) (forall (t :: [Maybe Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @sh')
      = IShX esh
-> Mixed (esh ++ sh) (Mixed sh' a)
-> Mixed esh (Mixed sh (Mixed sh' a))
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
IShX sh1 -> Mixed (sh1 ++ sh2) a -> Mixed sh1 (Mixed sh2 a)
M_Nest IShX esh
esh (IShX (esh ++ sh)
-> Mixed ((esh ++ sh) ++ sh') a -> Mixed (esh ++ sh) (Mixed sh' a)
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
IShX sh1 -> Mixed (sh1 ++ sh2) a -> Mixed sh1 (Mixed sh2 a)
M_Nest (Proxy sh'
-> StaticShX (esh ++ sh)
-> ShX ((esh ++ sh) ++ sh') Int
-> IShX (esh ++ sh)
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 @sh') (IShX esh -> StaticShX esh
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX IShX esh
esh StaticShX esh -> StaticShX sh -> StaticShX (esh ++ sh)
forall (sh :: [Maybe Nat]) (sh' :: [Maybe Nat]).
StaticShX sh -> StaticShX sh' -> StaticShX (sh ++ sh')
`ssxAppend` StaticShX sh
ssh) (Mixed ((esh ++ sh) ++ sh') a -> ShX ((esh ++ sh) ++ sh') Int
forall (sh :: [Maybe Nat]). Mixed sh a -> IShX sh
forall a (sh :: [Maybe Nat]). Elt a => Mixed sh a -> IShX sh
mshape Mixed (esh ++ (sh ++ sh')) a
Mixed ((esh ++ sh) ++ sh') a
x)) Mixed (esh ++ (sh ++ sh')) a
Mixed ((esh ++ sh) ++ sh') a
x)
    go (ConvUnnest @sh @sh') (M_Nest @esh IShX esh
esh (M_Nest IShX (esh ++ sh)
_ Mixed ((esh ++ sh) ++ sh') a
x))
      | ((esh ++ sh) ++ sh') :~: (esh ++ (sh ++ sh'))
Refl <- Proxy esh
-> Proxy sh
-> Proxy sh'
-> ((esh ++ sh) ++ sh') :~: (esh ++ (sh ++ sh'))
forall {a1} (a2 :: [a1]) (b :: [a1]) (c :: [a1]).
Proxy a2
-> Proxy b -> Proxy c -> ((a2 ++ b) ++ c) :~: (a2 ++ (b ++ c))
lemAppAssoc (forall (t :: [Maybe Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @esh) (forall (t :: [Maybe Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @sh) (forall (t :: [Maybe Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @sh')
      = IShX esh
-> Mixed (esh ++ (sh ++ sh')) a -> Mixed esh (Mixed (sh ++ sh') a)
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
IShX sh1 -> Mixed (sh1 ++ sh2) a -> Mixed sh1 (Mixed sh2 a)
M_Nest IShX esh
esh Mixed (esh ++ (sh ++ sh')) a
Mixed ((esh ++ sh) ++ sh') a
x
    go Conversion a b
ConvZip Mixed esh a
x =
      -- no need to check that the two esh's are equal because they were zipped previously
      let (M_Nest IShX esh
esh Mixed (esh ++ sh) a
x1, M_Nest IShX esh
_ Mixed (esh ++ sh) b
x2) = Mixed esh (Mixed sh a, Mixed sh b)
-> (Mixed esh (Mixed sh a), Mixed esh (Mixed sh b))
forall (sh :: [Maybe Nat]) a b.
Mixed sh (a, b) -> (Mixed sh a, Mixed sh b)
munzip Mixed esh a
Mixed esh (Mixed sh a, Mixed sh b)
x
      in IShX esh -> Mixed (esh ++ sh) (a, b) -> Mixed esh (Mixed sh (a, b))
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
IShX sh1 -> Mixed (sh1 ++ sh2) a -> Mixed sh1 (Mixed sh2 a)
M_Nest IShX esh
esh (Mixed (esh ++ sh) a
-> Mixed (esh ++ sh) b -> Mixed (esh ++ sh) (a, b)
forall a b (sh :: [Maybe Nat]).
(Elt a, Elt b) =>
Mixed sh a -> Mixed sh b -> Mixed sh (a, b)
mzip Mixed (esh ++ sh) a
x1 Mixed (esh ++ sh) b
x2)
    go Conversion a b
ConvUnzip (M_Nest IShX esh
esh Mixed (esh ++ sh) (a, b)
x) =
      let (Mixed (esh ++ sh) a
x1, Mixed (esh ++ sh) b
x2) = Mixed (esh ++ sh) (a, b)
-> (Mixed (esh ++ sh) a, Mixed (esh ++ sh) b)
forall (sh :: [Maybe Nat]) a b.
Mixed sh (a, b) -> (Mixed sh a, Mixed sh b)
munzip Mixed (esh ++ sh) (a, b)
x
      in Mixed esh (Mixed sh a)
-> Mixed esh (Mixed sh b) -> Mixed esh (Mixed sh a, Mixed sh b)
forall a b (sh :: [Maybe Nat]).
(Elt a, Elt b) =>
Mixed sh a -> Mixed sh b -> Mixed sh (a, b)
mzip (IShX esh -> Mixed (esh ++ sh) a -> Mixed esh (Mixed sh a)
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
IShX sh1 -> Mixed (sh1 ++ sh2) a -> Mixed sh1 (Mixed sh2 a)
M_Nest IShX esh
esh Mixed (esh ++ sh) a
x1) (IShX esh -> Mixed (esh ++ sh) b -> Mixed esh (Mixed sh b)
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
IShX sh1 -> Mixed (sh1 ++ sh2) a -> Mixed sh1 (Mixed sh2 a)
M_Nest IShX esh
esh Mixed (esh ++ sh) b
x2)

    lemRankAppRankEq :: Rank sh ~ Rank sh'
                     => Proxy esh -> Proxy sh -> Proxy sh'
                     -> Rank (esh ++ sh) :~: Rank (esh ++ sh')
    lemRankAppRankEq :: forall {a} (sh :: [a]) (sh' :: [a]) (esh :: [a]).
(Rank sh ~ Rank sh') =>
Proxy esh
-> Proxy sh -> Proxy sh' -> Rank (esh ++ sh) :~: Rank (esh ++ sh')
lemRankAppRankEq Proxy esh
_ Proxy sh
_ Proxy sh'
_ = Rank (esh ++ sh) :~: Rank (esh ++ sh')
forall {k} (a :: k) (b :: k). a :~: b
unsafeCoerceRefl

    lemRankAppRankEqRepNo :: Proxy esh -> Proxy sh
                          -> Rank (esh ++ sh) :~: Rank (esh ++ Replicate (Rank sh) Nothing)
    lemRankAppRankEqRepNo :: forall {a} (esh :: [Maybe a]) (sh :: [Maybe a]).
Proxy esh
-> Proxy sh
-> Rank (esh ++ sh) :~: Rank (esh ++ Replicate (Rank sh) 'Nothing)
lemRankAppRankEqRepNo Proxy esh
_ Proxy sh
_ = Rank (esh ++ sh) :~: Rank (esh ++ Replicate (Rank sh) 'Nothing)
forall {k} (a :: k) (b :: k). a :~: b
unsafeCoerceRefl

    lemRankAppRankEqMapJust :: Rank sh ~ Rank sh'
                            => Proxy esh -> Proxy sh -> Proxy sh'
                            -> Rank (esh ++ sh) :~: Rank (esh ++ MapJust sh')
    lemRankAppRankEqMapJust :: forall {a} (sh :: [Maybe a]) (sh' :: [a]) (esh :: [Maybe a]).
(Rank sh ~ Rank sh') =>
Proxy esh
-> Proxy sh
-> Proxy sh'
-> Rank (esh ++ sh) :~: Rank (esh ++ MapJust sh')
lemRankAppRankEqMapJust Proxy esh
_ Proxy sh
_ Proxy sh'
_ = Rank (esh ++ sh) :~: Rank (esh ++ MapJust sh')
forall {k} (a :: k) (b :: k). a :~: b
unsafeCoerceRefl


-- * Special cases of array conversions

mcast :: forall sh1 sh2 a. (Rank sh1 ~ Rank sh2, Elt a)
      => StaticShX sh2 -> Mixed sh1 a -> Mixed sh2 a
mcast :: forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
(Rank sh1 ~ Rank sh2, Elt a) =>
StaticShX sh2 -> Mixed sh1 a -> Mixed sh2 a
mcast StaticShX sh2
ssh2 Mixed sh1 a
arr
  | (sh1 ++ '[]) :~: sh1
Refl <- forall (l :: [Maybe Nat]). (l ++ '[]) :~: l
forall {a} (l :: [a]). (l ++ '[]) :~: l
lemAppNil @sh1
  , (sh2 ++ '[]) :~: sh2
Refl <- forall (l :: [Maybe Nat]). (l ++ '[]) :~: l
forall {a} (l :: [a]). (l ++ '[]) :~: l
lemAppNil @sh2
  = StaticShX sh1
-> StaticShX sh2
-> Proxy '[]
-> Mixed (sh1 ++ '[]) a
-> Mixed (sh2 ++ '[]) 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 (ShX sh1 Int -> StaticShX sh1
forall (sh :: [Maybe Nat]) i. ShX sh i -> StaticShX sh
ssxFromShX (Mixed sh1 a -> ShX sh1 Int
forall (sh :: [Maybe Nat]). Mixed sh a -> IShX sh
forall a (sh :: [Maybe Nat]). Elt a => Mixed sh a -> IShX sh
mshape Mixed sh1 a
arr)) StaticShX sh2
ssh2 (forall (t :: [Maybe Nat]). Proxy t
forall {k} (t :: k). Proxy t
Proxy @'[]) Mixed sh1 a
Mixed (sh1 ++ '[]) a
arr

mtoRanked :: forall sh a. Elt a => Mixed sh a -> Ranked (Rank sh) a
mtoRanked :: forall (sh :: [Maybe Nat]) a.
Elt a =>
Mixed sh a -> Ranked (Rank sh) a
mtoRanked = Conversion (Mixed sh a) (Ranked (Rank sh) a)
-> Mixed sh a -> Ranked (Rank sh) a
forall a b. (Elt a, Elt b) => Conversion a b -> a -> b
convert Conversion (Mixed sh a) (Ranked (Rank sh) a)
forall sh (sh' :: [Maybe Nat]).
Elt sh =>
Conversion (Mixed sh' sh) (Ranked (Rank sh') sh)
ConvXR

rtoMixed :: forall n a. Ranked n a -> Mixed (Replicate n Nothing) a
rtoMixed :: forall (n :: Nat) a. Ranked n a -> Mixed (Replicate n 'Nothing) a
rtoMixed (Ranked Mixed (Replicate n 'Nothing) a
arr) = Mixed (Replicate n 'Nothing) a
arr

-- | A more weakly-typed version of 'rtoMixed' that does a runtime shape
-- compatibility check.
rcastToMixed :: (Rank sh ~ n, Elt a) => StaticShX sh -> Ranked n a -> Mixed sh a
rcastToMixed :: forall (sh :: [Maybe Nat]) (n :: Nat) a.
(Rank sh ~ n, Elt a) =>
StaticShX sh -> Ranked n a -> Mixed sh a
rcastToMixed StaticShX sh
sshx rarr :: Ranked n a
rarr@(Ranked Mixed (Replicate n 'Nothing) a
arr)
  | Rank (Replicate n 'Nothing) :~: n
Refl <- SNat n -> Rank (Replicate n 'Nothing) :~: n
forall (proxy :: Nat -> *) (n :: Nat).
proxy n -> Rank (Replicate n 'Nothing) :~: n
lemRankReplicate (Ranked n a -> SNat n
forall a (n :: Nat). Elt a => Ranked n a -> SNat n
rrank Ranked n a
rarr)
  = StaticShX sh -> Mixed (Replicate n 'Nothing) a -> Mixed sh a
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
(Rank sh1 ~ Rank sh2, Elt a) =>
StaticShX sh2 -> Mixed sh1 a -> Mixed sh2 a
mcast StaticShX sh
sshx Mixed (Replicate n 'Nothing) a
arr

mcastToShaped :: forall sh sh' a. (Elt a, Rank sh ~ Rank sh')
              => ShS sh' -> Mixed sh a -> Shaped sh' a
mcastToShaped :: forall (sh :: [Maybe Nat]) (sh' :: [Nat]) a.
(Elt a, Rank sh ~ Rank sh') =>
ShS sh' -> Mixed sh a -> Shaped sh' a
mcastToShaped ShS sh'
targetsh = Conversion (Mixed sh a) (Shaped sh' a)
-> Mixed sh a -> Shaped sh' a
forall a b. (Elt a, Elt b) => Conversion a b -> a -> b
convert (ShS sh' -> Conversion (Mixed sh a) (Shaped sh' a)
forall (sh :: [Maybe Nat]) (sh' :: [Nat]) a.
(Rank sh ~ Rank sh', Elt a) =>
ShS sh' -> Conversion (Mixed sh a) (Shaped sh' a)
ConvXS' ShS sh'
targetsh)

stoMixed :: forall sh a. Shaped sh a -> Mixed (MapJust sh) a
stoMixed :: forall (sh :: [Nat]) a. Shaped sh a -> Mixed (MapJust sh) a
stoMixed (Shaped Mixed (MapJust sh) a
arr) = Mixed (MapJust sh) a
arr

-- | A more weakly-typed version of 'stoMixed' that does a runtime shape
-- compatibility check.
scastToMixed :: forall sh sh' a. (Elt a, Rank sh ~ Rank sh')
             => StaticShX sh' -> Shaped sh a -> Mixed sh' a
scastToMixed :: forall (sh :: [Nat]) (sh' :: [Maybe Nat]) a.
(Elt a, Rank sh ~ Rank sh') =>
StaticShX sh' -> Shaped sh a -> Mixed sh' a
scastToMixed StaticShX sh'
sshx 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)
  = StaticShX sh' -> Mixed (MapJust sh) a -> Mixed sh' a
forall (sh1 :: [Maybe Nat]) (sh2 :: [Maybe Nat]) a.
(Rank sh1 ~ Rank sh2, Elt a) =>
StaticShX sh2 -> Mixed sh1 a -> Mixed sh2 a
mcast StaticShX sh'
sshx Mixed (MapJust sh) a
arr

stoRanked :: Elt a => Shaped sh a -> Ranked (Rank sh) a
stoRanked :: forall a (sh :: [Nat]). Elt a => Shaped sh a -> Ranked (Rank sh) a
stoRanked 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)
  = Mixed (MapJust sh) a -> Ranked (Rank (MapJust sh)) a
forall (sh :: [Maybe Nat]) a.
Elt a =>
Mixed sh a -> Ranked (Rank sh) a
mtoRanked Mixed (MapJust sh) a
arr

rcastToShaped :: Elt a => Ranked (Rank sh) a -> ShS sh -> Shaped sh a
rcastToShaped :: forall a (sh :: [Nat]).
Elt a =>
Ranked (Rank sh) a -> ShS sh -> Shaped sh a
rcastToShaped (Ranked Mixed (Replicate (Rank sh) 'Nothing) a
arr) ShS sh
targetsh
  | Rank (Replicate (Rank (MapJust sh)) 'Nothing) :~: Rank (MapJust sh)
Refl <- SNat (Rank (MapJust sh))
-> Rank (Replicate (Rank (MapJust sh)) 'Nothing)
   :~: Rank (MapJust sh)
forall (proxy :: Nat -> *) (n :: Nat).
proxy n -> Rank (Replicate n 'Nothing) :~: n
lemRankReplicate (ShX (MapJust sh) Int -> SNat (Rank (MapJust sh))
forall (sh :: [Maybe Nat]) i. ShX sh i -> SNat (Rank sh)
shxRank (ShS sh -> ShX (MapJust sh) Int
forall (sh :: [Nat]). ShS sh -> IShX (MapJust sh)
shxFromShS ShS sh
targetsh))
  , Rank (MapJust sh) :~: Rank sh
Refl <- ShS sh -> Rank (MapJust sh) :~: Rank sh
forall (sh :: [Nat]). ShS sh -> Rank (MapJust sh) :~: Rank sh
lemRankMapJust ShS sh
targetsh
  = ShS sh -> Mixed (Replicate (Rank sh) 'Nothing) a -> Shaped sh a
forall (sh :: [Maybe Nat]) (sh' :: [Nat]) a.
(Elt a, Rank sh ~ Rank sh') =>
ShS sh' -> Mixed sh a -> Shaped sh' a
mcastToShaped ShS sh
targetsh Mixed (Replicate (Rank sh) 'Nothing) a
arr