path: root/src/Data/Array/Nested/Internal.hs

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RoleAnnotations #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE StandaloneKindSignatures #-}
{-# LANGUAGE TypeAbstractions #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}

{-# OPTIONS -Wno-unused-imports #-}

{-|
TODO:
* Allow downtyping certain dimensions, and write conversions between Mixed,
  Ranked and Shaped

* Mikolaj wants these:

    About your wishlist of operations: these are already there

      OR.index
      OR.append
      OR.transpose

    These can be easily lifted from the definition for XArray (5min work):

      OR.scalar
      OR.unScalar
      OR.constant

    These should not be hard:

      OR.fromList
      ORB.toList . OR.unravel
      OR.ravel . ORB.fromList
      OR.slice
      OR.rev
      OR.reshape

    though it's a bit unfortunate that we end up needing toList. Looking in
    horde-ad I see that you seem to need them to do certain operations in Haskell
    that orthotope doesn't support?

    For this one we'll need to see to what extent you really need it, and what API
    you'd need precisely:

      OR.rerank

    and for these we have an API design question:

      OR.toVector
      OR.fromVector

-}

module Data.Array.Nested.Internal where

import Prelude hiding (mappend)

import Control.Monad (forM_, when)
import Control.Monad.ST
import qualified Data.Array.RankedS as S
import Data.Bifunctor (first)
import Data.Coerce (coerce, Coercible)
import Data.Foldable (toList)
import Data.Functor.Const
import Data.Kind
import Data.List.NonEmpty (NonEmpty(..))
import Data.Proxy
import Data.Type.Equality
import qualified Data.Vector.Storable as VS
import qualified Data.Vector.Storable.Mutable as VSM
import Foreign.Storable (Storable)
import GHC.TypeLits
import qualified GHC.TypeNats as TypeNats
import Unsafe.Coerce

import Data.Array.Mixed
import qualified Data.Array.Mixed as X


-- Invariant in the API
-- ====================
--
-- In the underlying XArray, there is some shape for elements of an empty
-- array. For example, for this array:
--
--   arr :: Ranked I3 (Ranked I2 Int, Ranked I1 Float)
--   rshape arr == 0 :.: 0 :.: 0 :.: ZIR
--
-- the two underlying XArrays have a shape, and those shapes might be anything.
-- The invariant is that these element shapes are unobservable in the API.
-- (This is possible because you ought to not be able to get to such an element
-- without indexing out of bounds.)
--
-- Note, though, that the converse situation may arise: the outer array might
-- be nonempty but then the inner arrays might. This is fine, an invariant only
-- applies if the _outer_ array is empty.
--
-- TODO: can we enforce that the elements of an empty (nested) array have
-- all-zero shape?
--   -> no, because mlift and also any kind of internals probing from outsiders


-- Primitive element types
-- =======================
--
-- There are a few primitive element types; arrays containing elements of such
-- type are a newtype over an XArray, which it itself a newtype over a Vector.
-- Unfortunately, the setup of the library requires us to list these primitive
-- element types multiple times; to aid in extending the list, all these lists
-- have been marked with [PRIMITIVE ELEMENT TYPES LIST].


type family MapJust l where
  MapJust '[] = '[]
  MapJust (x : xs) = Just x : MapJust xs


-- Stupid things that the type checker should be able to figure out in-line, but can't

subst1 :: forall f a b. a :~: b -> f a :~: f b
subst1 Refl = Refl

subst2 :: forall f c a b. a :~: b -> f a c :~: f b c
subst2 Refl = Refl

lemAppLeft :: Proxy l -> a :~: b -> a ++ l :~: b ++ l
lemAppLeft _ Refl = Refl

knownNatSucc :: KnownNat n => Dict KnownNat (n + 1)
knownNatSucc = Dict


lemKnownShX :: StaticShX sh -> Dict KnownShX sh
lemKnownShX ZKX = Dict
lemKnownShX (SKnown GHC_SNat :!% ssh) | Dict <- lemKnownShX ssh = Dict
lemKnownShX (SUnknown () :!% ssh) | Dict <- lemKnownShX ssh = Dict

ssxFromSNat :: SNat n -> StaticShX (Replicate n Nothing)
ssxFromSNat SZ = ZKX
ssxFromSNat (SS (n :: SNat nm1)) | Refl <- X.lemReplicateSucc @(Nothing @Nat) @nm1 = SUnknown () :!% ssxFromSNat n

lemKnownReplicate :: SNat n -> Dict KnownShX (Replicate n Nothing)
lemKnownReplicate sn = lemKnownShX (ssxFromSNat sn)

lemRankReplicate :: SNat n -> X.Rank (Replicate n (Nothing @Nat)) :~: n
lemRankReplicate SZ = Refl
lemRankReplicate (SS (n :: SNat nm1))
  | Refl <- X.lemReplicateSucc @(Nothing @Nat) @nm1
  , Refl <- lemRankReplicate n
  = Refl

lemRankMapJust :: forall sh. ShS sh -> X.Rank (MapJust sh) :~: X.Rank sh
lemRankMapJust ZSS = Refl
lemRankMapJust (_ :$$ sh') | Refl <- lemRankMapJust sh' = Refl

lemReplicatePlusApp :: forall n m a. SNat n -> Proxy m -> Proxy a
                    -> Replicate (n + m) a :~: Replicate n a ++ Replicate m a
lemReplicatePlusApp sn _ _ = go sn
  where
    go :: SNat n' -> Replicate (n' + m) a :~: Replicate n' a ++ Replicate m a
    go SZ = Refl
    go (SS (n :: SNat n'm1))
      | Refl <- X.lemReplicateSucc @a @n'm1
      , Refl <- go n
      = sym (X.lemReplicateSucc @a @(n'm1 + m))


-- === NEW INDEX TYPES === --

type role ListR nominal representational
type ListR :: Nat -> Type -> Type
data ListR n i where
  ZR :: ListR 0 i
  (:::) :: forall n {i}. i -> ListR n i -> ListR (n + 1) i
deriving instance Show i => Show (ListR n i)
deriving instance Eq i => Eq (ListR n i)
deriving instance Ord i => Ord (ListR n i)
deriving instance Functor (ListR n)
deriving instance Foldable (ListR n)
infixr 3 :::

data UnconsListRRes i n1 =
  forall n. (n + 1 ~ n1) => UnconsListRRes (ListR n i) i
unconsListR :: ListR n1 i -> Maybe (UnconsListRRes i n1)
unconsListR (i ::: sh') = Just (UnconsListRRes sh' i)
unconsListR ZR = Nothing


-- | An index into a rank-typed array.
type role IxR nominal representational
type IxR :: Nat -> Type -> Type
newtype IxR n i = IxR (ListR n i)
  deriving (Show, Eq, Ord)
  deriving newtype (Functor, Foldable)

pattern ZIR :: forall n i. () => n ~ 0 => IxR n i
pattern ZIR = IxR ZR

pattern (:.:)
  :: forall {n1} {i}.
     forall n. (n + 1 ~ n1)
  => i -> IxR n i -> IxR n1 i
pattern i :.: sh <- IxR (unconsListR -> Just (UnconsListRRes (IxR -> sh) i))
  where i :.: IxR sh = IxR (i ::: sh)
infixr 3 :.:

{-# COMPLETE ZIR, (:.:) #-}

type IIxR n = IxR n Int


type role ShR nominal representational
type ShR :: Nat -> Type -> Type
newtype ShR n i = ShR (ListR n i)
  deriving (Show, Eq, Ord)
  deriving newtype (Functor, Foldable)

type IShR n = ShR n Int

pattern ZSR :: forall n i. () => n ~ 0 => ShR n i
pattern ZSR = ShR ZR

pattern (:$:)
  :: forall {n1} {i}.
     forall n. (n + 1 ~ n1)
  => i -> ShR n i -> ShR n1 i
pattern i :$: sh <- ShR (unconsListR -> Just (UnconsListRRes (ShR -> sh) i))
  where i :$: (ShR sh) = ShR (i ::: sh)
infixr 3 :$:

{-# COMPLETE ZSR, (:$:) #-}


type role ListS nominal representational
type ListS :: [Nat] -> (Nat -> Type) -> Type
data ListS sh f where
  ZS :: ListS '[] f
  (::$) :: forall n sh {f}. f n -> ListS sh f -> ListS (n : sh) f
deriving instance (forall n. Show (f n)) => Show (ListS sh f)
deriving instance (forall n. Eq (f n)) => Eq (ListS sh f)
deriving instance (forall n. Ord (f n)) => Ord (ListS sh f)
infixr 3 ::$

data UnconsListSRes f sh1 =
  forall n sh. (n : sh ~ sh1) => UnconsListSRes (ListS sh f) (f n)
unconsListS :: ListS sh1 f -> Maybe (UnconsListSRes f sh1)
unconsListS (x ::$ sh') = Just (UnconsListSRes sh' x)
unconsListS ZS = Nothing

fmapListS :: (forall n. f n -> g n) -> ListS sh f -> ListS sh g
fmapListS _ ZS = ZS
fmapListS f (x ::$ xs) = f x ::$ fmapListS f xs

foldListS :: Monoid m => (forall n. f n -> m) -> ListS sh f -> m
foldListS _ ZS = mempty
foldListS f (x ::$ xs) = f x <> foldListS f xs


-- | An index into a shape-typed array.
--
-- For convenience, this contains regular 'Int's instead of bounded integers
-- (traditionally called \"@Fin@\"). Note that because the shape of a
-- shape-typed array is known statically, you can also retrieve the array shape
-- from a 'KnownShape' dictionary.
type role IxS nominal representational
type IxS :: [Nat] -> Type -> Type
newtype IxS sh i = IxS (ListS sh (Const i))
  deriving (Show, Eq, Ord)

pattern ZIS :: forall sh i. () => sh ~ '[] => IxS sh i
pattern ZIS = IxS ZS

pattern (:.$)
  :: forall {sh1} {i}.
     forall n sh. (n : sh ~ sh1)
  => i -> IxS sh i -> IxS sh1 i
pattern i :.$ shl <- IxS (unconsListS -> Just (UnconsListSRes (IxS -> shl) (getConst -> i)))
  where i :.$ IxS shl = IxS (Const i ::$ shl)
infixr 3 :.$

{-# COMPLETE ZIS, (:.$) #-}

type IIxS sh = IxS sh Int

instance Functor (IxS sh) where
  fmap f (IxS l) = IxS (fmapListS (Const . f . getConst) l)

instance Foldable (IxS sh) where
  foldMap f (IxS l) = foldListS (f . getConst) l


-- | The shape of a shape-typed array given as a list of 'SNat' values.
type role ShS nominal
type ShS :: [Nat] -> Type
newtype ShS sh = ShS (ListS sh SNat)
  deriving (Show, Eq, Ord)

pattern ZSS :: forall sh. () => sh ~ '[] => ShS sh
pattern ZSS = ShS ZS

pattern (:$$)
  :: forall {sh1}.
     forall n sh. (n : sh ~ sh1)
  => SNat n -> ShS sh -> ShS sh1
pattern i :$$ shl <- ShS (unconsListS -> Just (UnconsListSRes (ShS -> shl) i))
  where i :$$ ShS shl = ShS (i ::$ shl)

infixr 3 :$$

{-# COMPLETE ZSS, (:$$) #-}


-- | Wrapper type used as a tag to attach instances on. The instances on arrays
-- of @'Primitive' a@ are more polymorphic than the direct instances for arrays
-- of scalars; this means that if @orthotope@ supports an element type @T@ that
-- this library does not (directly), it may just work if you use an array of
-- @'Primitive' T@ instead.
newtype Primitive a = Primitive a

-- | Element types that are primitive; arrays of these types are just a newtype
-- wrapper over an array.
class PrimElt a where
  fromPrimitive :: Mixed sh (Primitive a) -> Mixed sh a
  toPrimitive :: Mixed sh a -> Mixed sh (Primitive a)

  default fromPrimitive :: Coercible (Mixed sh a) (Mixed sh (Primitive a)) => Mixed sh (Primitive a) -> Mixed sh a
  fromPrimitive = coerce

  default toPrimitive :: Coercible (Mixed sh (Primitive a)) (Mixed sh a) => Mixed sh a -> Mixed sh (Primitive a)
  toPrimitive = coerce

-- [PRIMITIVE ELEMENT TYPES LIST]
instance PrimElt Int
instance PrimElt Double
instance PrimElt ()


-- | Mixed arrays: some dimensions are size-typed, some are not. Distributes
-- over product-typed elements using a data family so that the full array is
-- always in struct-of-arrays format.
--
-- Built on top of 'XArray' which is built on top of @orthotope@, meaning that
-- dimension permutations (e.g. 'mtranspose') are typically free.
--
-- Many of the methods for working on 'Mixed' arrays come from the 'Elt' type
-- class.
type Mixed :: [Maybe Nat] -> Type -> Type
data family Mixed sh a
-- NOTE: When opening up the Mixed abstraction, you might see dimension sizes
-- that you're not supposed to see. In particular, you might see (nonempty)
-- sizes of the elements of an empty array, which is information that should
-- ostensibly not exist; the full array is still empty.

data instance Mixed sh (Primitive a) = M_Primitive !(IShX sh) !(XArray sh a)
  deriving (Show)

-- [PRIMITIVE ELEMENT TYPES LIST]
newtype instance Mixed sh Int = M_Int (Mixed sh (Primitive Int))
  deriving (Show)
newtype instance Mixed sh Double = M_Double (Mixed sh (Primitive Double))
  deriving (Show)
newtype instance Mixed sh () = M_Nil (Mixed sh (Primitive ()))  -- no content, orthotope optimises this (via Vector)
  deriving (Show)
-- etc.

data instance Mixed sh (a, b) = M_Tup2 !(Mixed sh a) !(Mixed sh b)
deriving instance (Show (Mixed sh a), Show (Mixed sh b)) => Show (Mixed sh (a, b))
-- etc.

data instance Mixed sh1 (Mixed sh2 a) = M_Nest !(IShX sh1) !(Mixed (sh1 ++ sh2) a)
deriving instance Show (Mixed (sh1 ++ sh2) a) => Show (Mixed sh1 (Mixed sh2 a))


-- | Internal helper data family mirroring 'Mixed' that consists of mutable
-- vectors instead of 'XArray's.
type MixedVecs :: Type -> [Maybe Nat] -> Type -> Type
data family MixedVecs s sh a

newtype instance MixedVecs s sh (Primitive a) = MV_Primitive (VS.MVector s a)

-- [PRIMITIVE ELEMENT TYPES LIST]
newtype instance MixedVecs s sh Int = MV_Int (VS.MVector s Int)
newtype instance MixedVecs s sh Double = MV_Double (VS.MVector s Double)
newtype instance MixedVecs s sh () = MV_Nil (VS.MVector s ())  -- no content, MVector optimises this
-- etc.

data instance MixedVecs s sh (a, b) = MV_Tup2 !(MixedVecs s sh a) !(MixedVecs s sh b)
-- etc.

data instance MixedVecs s sh1 (Mixed sh2 a) = MV_Nest !(IShX sh2) !(MixedVecs s (sh1 ++ sh2) a)


-- | Tree giving the shape of every array component.
type family ShapeTree a where
  ShapeTree (Primitive _) = ()
  -- [PRIMITIVE ELEMENT TYPES LIST]
  ShapeTree Int = ()
  ShapeTree Double = ()
  ShapeTree () = ()

  ShapeTree (a, b) = (ShapeTree a, ShapeTree b)
  ShapeTree (Mixed sh a) = (IShX sh, ShapeTree a)
  ShapeTree (Ranked n a) = (IShR n, ShapeTree a)
  ShapeTree (Shaped sh a) = (ShS sh, ShapeTree a)


-- | Allowable element types in a mixed array, and by extension in a 'Ranked' or
-- 'Shaped' array. Note the polymorphic instance for 'Elt' of @'Primitive'
-- a@; see the documentation for 'Primitive' for more details.
class Elt a where
  -- ====== PUBLIC METHODS ====== --

  mshape :: Mixed sh a -> IShX sh
  mindex :: Mixed sh a -> IIxX sh -> a
  mindexPartial :: forall sh sh'. Mixed (sh ++ sh') a -> IIxX sh -> Mixed sh' a
  mscalar :: a -> Mixed '[] a

  -- | All arrays in the list, even subarrays inside @a@, must have the same
  -- shape; if they do not, a runtime error will be thrown. See the
  -- documentation of 'mgenerate' for more information about this restriction.
  -- Furthermore, the length of the list must correspond with @n@: if @n@ is
  -- @Just m@ and @m@ does not equal the length of the list, a runtime error is
  -- thrown.
  --
  -- If you want a single-dimensional array from your list, map 'mscalar'
  -- first.
  mfromList1 :: forall sh. NonEmpty (Mixed sh a) -> Mixed (Nothing : sh) a

  mtoList1 :: Mixed (n : sh) a -> [Mixed sh a]

  -- | Note: this library makes no particular guarantees about the shapes of
  -- arrays "inside" an empty array. With 'mlift' and 'mlift2' you can see the
  -- full 'XArray' and as such you can distinguish different empty arrays by
  -- the "shapes" of their elements. This information is meaningless, so you
  -- should not use it.
  mlift :: forall sh1 sh2.
           StaticShX sh2
        -> (forall sh' b. Storable b => StaticShX sh' -> XArray (sh1 ++ sh') b -> XArray (sh2 ++ sh') b)
        -> Mixed sh1 a -> Mixed sh2 a

  -- | See the documentation for 'mlift'.
  mlift2 :: forall sh1 sh2 sh3.
            StaticShX sh3
         -> (forall sh' 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

  mcast :: forall sh1 sh2 sh'. X.Rank sh1 ~ X.Rank sh2
        => StaticShX sh1 -> IShX sh2 -> Proxy sh' -> Mixed (sh1 ++ sh') a -> Mixed (sh2 ++ sh') a

  -- ====== PRIVATE METHODS ====== --

  mshapeTree :: a -> ShapeTree a

  mshapeTreeEq :: Proxy a -> ShapeTree a -> ShapeTree a -> Bool

  mshapeTreeEmpty :: Proxy a -> ShapeTree a -> Bool

  mshowShapeTree :: Proxy a -> ShapeTree a -> String

  -- | Given the shape of this array, an index and a value, write the value at
  -- that index in the vectors.
  mvecsWrite :: IShX sh -> IIxX sh -> a -> MixedVecs s sh a -> ST s ()

  -- | Given the shape of this array, an index and a value, write the value at
  -- that index in the vectors.
  mvecsWritePartial :: IShX (sh ++ sh') -> IIxX sh -> Mixed sh' a -> MixedVecs s (sh ++ sh') a -> ST s ()

  -- | Given the shape of this array, finalise the vectors into 'XArray's.
  mvecsFreeze :: IShX sh -> MixedVecs s sh a -> ST s (Mixed sh a)


-- | Element types for which we have evidence of the (static part of the) shape
-- in a type class constraint. Compare the instance contexts of the instances
-- of this class with those of 'Elt': some instances have an additional
-- "known-shape" constraint.
--
-- This class is (currently) only required for 'mgenerate' / 'rgenerate' /
-- 'sgenerate'.
class Elt a => KnownElt a where
  -- | Create an empty array. The given shape must have size zero; this may or may not be checked.
  memptyArray :: IShX sh -> Mixed sh a

  -- | Create uninitialised vectors for this array type, given the shape of
  -- this vector and an example for the contents.
  mvecsUnsafeNew :: IShX sh -> a -> ST s (MixedVecs s sh a)

  mvecsNewEmpty :: Proxy a -> ST s (MixedVecs s sh a)


-- Arrays of scalars are basically just arrays of scalars.
instance Storable a => Elt (Primitive a) where
  mshape (M_Primitive sh _) = sh
  mindex (M_Primitive _ a) i = Primitive (X.index a i)
  mindexPartial (M_Primitive sh a) i = M_Primitive (X.shDropIx sh i) (X.indexPartial a i)
  mscalar (Primitive x) = M_Primitive ZSX (X.scalar x)
  mfromList1 l@(arr1 :| _) =
    let sh = SUnknown (length l) :$% mshape arr1
    in M_Primitive sh (X.fromList1 (X.staticShapeFrom sh) (map (\(M_Primitive _ a) -> a) (toList l)))
  mtoList1 (M_Primitive sh arr) = map (M_Primitive (X.shTail sh)) (X.toList1 arr)

  mlift :: forall sh1 sh2.
           StaticShX sh2
        -> (StaticShX '[] -> XArray (sh1 ++ '[]) a -> XArray (sh2 ++ '[]) a)
        -> Mixed sh1 (Primitive a) -> Mixed sh2 (Primitive a)
  mlift ssh2 f (M_Primitive _ a)
    | Refl <- X.lemAppNil @sh1
    , Refl <- X.lemAppNil @sh2
    , let result = f ZKX a
    = M_Primitive (X.shape ssh2 result) result

  mlift2 :: forall sh1 sh2 sh3.
            StaticShX sh3
         -> (StaticShX '[] -> XArray (sh1 ++ '[]) a -> XArray (sh2 ++ '[]) a -> XArray (sh3 ++ '[]) a)
         -> Mixed sh1 (Primitive a) -> Mixed sh2 (Primitive a) -> Mixed sh3 (Primitive a)
  mlift2 ssh3 f (M_Primitive _ a) (M_Primitive _ b)
    | Refl <- X.lemAppNil @sh1
    , Refl <- X.lemAppNil @sh2
    , Refl <- X.lemAppNil @sh3
    , let result = f ZKX a b
    = M_Primitive (X.shape ssh3 result) result

  mcast :: forall sh1 sh2 sh'. X.Rank sh1 ~ X.Rank sh2
        => StaticShX sh1 -> IShX sh2 -> Proxy sh' -> Mixed (sh1 ++ sh') (Primitive a) -> Mixed (sh2 ++ sh') (Primitive a)
  mcast ssh1 sh2 _ (M_Primitive sh1' arr) =
    let (_, sh') = shAppSplit (Proxy @sh') ssh1 sh1'
    in M_Primitive (shAppend sh2 sh') (X.cast ssh1 sh2 (X.staticShapeFrom sh') arr)

  mshapeTree _ = ()
  mshapeTreeEq _ () () = True
  mshapeTreeEmpty _ () = False
  mshowShapeTree _ () = "()"
  mvecsWrite sh i (Primitive x) (MV_Primitive v) = VSM.write v (X.toLinearIdx sh i) x

  -- TODO: this use of toVector is suboptimal
  mvecsWritePartial
    :: forall sh' sh s.
       IShX (sh ++ sh') -> IIxX sh -> Mixed sh' (Primitive a) -> MixedVecs s (sh ++ sh') (Primitive a) -> ST s ()
  mvecsWritePartial sh i (M_Primitive sh' arr) (MV_Primitive v) = do
    let arrsh = X.shape (X.staticShapeFrom sh') arr
        offset = X.toLinearIdx sh (X.ixAppend i (X.zeroIxX' arrsh))
    VS.copy (VSM.slice offset (X.shapeSize arrsh) v) (X.toVector arr)

  mvecsFreeze sh (MV_Primitive v) = M_Primitive sh . X.fromVector sh <$> VS.freeze v

-- [PRIMITIVE ELEMENT TYPES LIST]
deriving via Primitive Int instance Elt Int
deriving via Primitive Double instance Elt Double
deriving via Primitive () instance Elt ()

instance Storable a => KnownElt (Primitive a) where
  memptyArray sh = M_Primitive sh (X.empty sh)
  mvecsUnsafeNew sh _ = MV_Primitive <$> VSM.unsafeNew (X.shapeSize sh)
  mvecsNewEmpty _ = MV_Primitive <$> VSM.unsafeNew 0

-- [PRIMITIVE ELEMENT TYPES LIST]
deriving via Primitive Int instance KnownElt Int
deriving via Primitive Double instance KnownElt Double
deriving via Primitive () instance KnownElt ()

-- Arrays of pairs are pairs of arrays.
instance (Elt a, Elt b) => Elt (a, b) where
  mshape (M_Tup2 a _) = mshape a
  mindex (M_Tup2 a b) i = (mindex a i, mindex b i)
  mindexPartial (M_Tup2 a b) i = M_Tup2 (mindexPartial a i) (mindexPartial b i)
  mscalar (x, y) = M_Tup2 (mscalar x) (mscalar y)
  mfromList1 l =
    M_Tup2 (mfromList1 ((\(M_Tup2 x _) -> x) <$> l))
           (mfromList1 ((\(M_Tup2 _ y) -> y) <$> l))
  mtoList1 (M_Tup2 a b) = zipWith M_Tup2 (mtoList1 a) (mtoList1 b)
  mlift ssh2 f (M_Tup2 a b) = M_Tup2 (mlift ssh2 f a) (mlift ssh2 f b)
  mlift2 ssh3 f (M_Tup2 a b) (M_Tup2 x y) = M_Tup2 (mlift2 ssh3 f a x) (mlift2 ssh3 f b y)

  mcast ssh1 sh2 psh' (M_Tup2 a b) =
    M_Tup2 (mcast ssh1 sh2 psh' a) (mcast ssh1 sh2 psh' b)

  mshapeTree (x, y) = (mshapeTree x, mshapeTree y)
  mshapeTreeEq _ (t1, t2) (t1', t2') = mshapeTreeEq (Proxy @a) t1 t1' && mshapeTreeEq (Proxy @b) t2 t2'
  mshapeTreeEmpty _ (t1, t2) = mshapeTreeEmpty (Proxy @a) t1 && mshapeTreeEmpty (Proxy @b) t2
  mshowShapeTree _ (t1, t2) = "(" ++ mshowShapeTree (Proxy @a) t1 ++ ", " ++ mshowShapeTree (Proxy @b) t2 ++ ")"
  mvecsWrite sh i (x, y) (MV_Tup2 a b) = do
    mvecsWrite sh i x a
    mvecsWrite sh i y b
  mvecsWritePartial sh i (M_Tup2 x y) (MV_Tup2 a b) = do
    mvecsWritePartial sh i x a
    mvecsWritePartial sh i y b
  mvecsFreeze sh (MV_Tup2 a b) = M_Tup2 <$> mvecsFreeze sh a <*> mvecsFreeze sh b

instance (KnownElt a, KnownElt b) => KnownElt (a, b) where
  memptyArray sh = M_Tup2 (memptyArray sh) (memptyArray sh)
  mvecsUnsafeNew sh (x, y) = MV_Tup2 <$> mvecsUnsafeNew sh x <*> mvecsUnsafeNew sh y
  mvecsNewEmpty _ = MV_Tup2 <$> mvecsNewEmpty (Proxy @a) <*> mvecsNewEmpty (Proxy @b)

-- Arrays of arrays are just arrays, but with more dimensions.
instance Elt a => Elt (Mixed sh' a) where
  -- TODO: this is quadratic in the nesting depth because it repeatedly
  -- truncates the shape vector to one a little shorter. Fix with a
  -- moverlongShape method, a prefix of which is mshape.
  mshape :: forall sh. Mixed sh (Mixed sh' a) -> IShX sh
  mshape (M_Nest sh arr)
    = fst (shAppSplit (Proxy @sh') (X.staticShapeFrom sh) (mshape arr))

  mindex :: Mixed sh (Mixed sh' a) -> IIxX sh -> Mixed sh' a
  mindex (M_Nest _ arr) i = mindexPartial arr i

  mindexPartial :: forall sh1 sh2.
                   Mixed (sh1 ++ sh2) (Mixed sh' a) -> IIxX sh1 -> Mixed sh2 (Mixed sh' a)
  mindexPartial (M_Nest sh arr) i
    | Refl <- X.lemAppAssoc (Proxy @sh1) (Proxy @sh2) (Proxy @sh')
    = M_Nest (X.shDropIx sh i) (mindexPartial @a @sh1 @(sh2 ++ sh') arr i)

  mscalar = M_Nest ZSX

  mfromList1 :: forall sh. NonEmpty (Mixed sh (Mixed sh' a)) -> Mixed (Nothing : sh) (Mixed sh' a)
  mfromList1 l@(arr :| _) =
    M_Nest (SUnknown (length l) :$% mshape arr)
           (mfromList1 ((\(M_Nest _ a) -> a) <$> l))

  mtoList1 (M_Nest sh arr) = map (M_Nest (X.shTail sh)) (mtoList1 arr)

  mlift :: forall sh1 sh2.
           StaticShX sh2
        -> (forall shT b. Storable b => StaticShX shT -> XArray (sh1 ++ shT) b -> XArray (sh2 ++ shT) b)
        -> Mixed sh1 (Mixed sh' a) -> Mixed sh2 (Mixed sh' a)
  mlift ssh2 f (M_Nest sh1 arr) =
    let result = mlift (X.ssxAppend ssh2 ssh') f' arr
        (sh2, _) = shAppSplit (Proxy @sh') ssh2 (mshape result)
    in M_Nest sh2 result
    where
      ssh' = X.staticShapeFrom (snd (shAppSplit (Proxy @sh') (X.staticShapeFrom sh1) (mshape arr)))

      f' :: forall shT b. Storable b => StaticShX shT -> XArray ((sh1 ++ sh') ++ shT) b -> XArray ((sh2 ++ sh') ++ shT) b
      f' sshT
        | Refl <- X.lemAppAssoc (Proxy @sh1) (Proxy @sh') (Proxy @shT)
        , Refl <- X.lemAppAssoc (Proxy @sh2) (Proxy @sh') (Proxy @shT)
        = f (X.ssxAppend ssh' sshT)

  mlift2 :: forall sh1 sh2 sh3.
            StaticShX sh3
         -> (forall shT b. Storable b => StaticShX shT -> XArray (sh1 ++ shT) b -> XArray (sh2 ++ shT) b -> XArray (sh3 ++ shT) b)
         -> Mixed sh1 (Mixed sh' a) -> Mixed sh2 (Mixed sh' a) -> Mixed sh3 (Mixed sh' a)
  mlift2 ssh3 f (M_Nest sh1 arr1) (M_Nest _ arr2) =
    let result = mlift2 (X.ssxAppend ssh3 ssh') f' arr1 arr2
        (sh3, _) = shAppSplit (Proxy @sh') ssh3 (mshape result)
    in M_Nest sh3 result
    where
      ssh' = X.staticShapeFrom (snd (shAppSplit (Proxy @sh') (X.staticShapeFrom sh1) (mshape arr1)))

      f' :: forall shT b. Storable b => StaticShX shT -> XArray ((sh1 ++ sh') ++ shT) b -> XArray ((sh2 ++ sh') ++ shT) b -> XArray ((sh3 ++ sh') ++ shT) b
      f' sshT
        | Refl <- X.lemAppAssoc (Proxy @sh1) (Proxy @sh') (Proxy @shT)
        , Refl <- X.lemAppAssoc (Proxy @sh2) (Proxy @sh') (Proxy @shT)
        , Refl <- X.lemAppAssoc (Proxy @sh3) (Proxy @sh') (Proxy @shT)
        = f (X.ssxAppend ssh' sshT)

  mcast :: forall sh1 sh2 shT. X.Rank sh1 ~ X.Rank sh2
        => StaticShX sh1 -> IShX sh2 -> Proxy shT -> Mixed (sh1 ++ shT) (Mixed sh' a) -> Mixed (sh2 ++ shT) (Mixed sh' a)
  mcast ssh1 sh2 _ (M_Nest sh1T arr)
    | Refl <- X.lemAppAssoc (Proxy @sh1) (Proxy @shT) (Proxy @sh')
    , Refl <- X.lemAppAssoc (Proxy @sh2) (Proxy @shT) (Proxy @sh')
    = let (_, shT) = shAppSplit (Proxy @shT) ssh1 sh1T
      in M_Nest (shAppend sh2 shT) (mcast ssh1 sh2 (Proxy @(shT ++ sh')) arr)

  mshapeTree :: Mixed sh' a -> ShapeTree (Mixed sh' a)
  mshapeTree arr = (mshape arr, mshapeTree (mindex arr (X.zeroIxX (X.staticShapeFrom (mshape arr)))))

  mshapeTreeEq _ (sh1, t1) (sh2, t2) = sh1 == sh2 && mshapeTreeEq (Proxy @a) t1 t2

  mshapeTreeEmpty _ (sh, t) = X.shapeSize sh == 0 && mshapeTreeEmpty (Proxy @a) t

  mshowShapeTree _ (sh, t) = "(" ++ show sh ++ ", " ++ mshowShapeTree (Proxy @a) t ++ ")"

  mvecsWrite sh idx val (MV_Nest sh' vecs) = mvecsWritePartial (X.shAppend sh sh') idx val vecs

  mvecsWritePartial :: forall sh1 sh2 s.
                       IShX (sh1 ++ sh2) -> IIxX sh1 -> Mixed sh2 (Mixed sh' a)
                    -> MixedVecs s (sh1 ++ sh2) (Mixed sh' a)
                    -> ST s ()
  mvecsWritePartial sh12 idx (M_Nest _ arr) (MV_Nest sh' vecs)
    | Refl <- X.lemAppAssoc (Proxy @sh1) (Proxy @sh2) (Proxy @sh')
    = mvecsWritePartial (X.shAppend sh12 sh') idx arr vecs

  mvecsFreeze sh (MV_Nest sh' vecs) = M_Nest sh <$> mvecsFreeze (X.shAppend sh sh') vecs

-- | Evidence for the static part of a shape. This pops up only when you are
-- polymorphic in the element type of an array.
type KnownShX :: [Maybe Nat] -> Constraint
class KnownShX sh where knownShX :: StaticShX sh
instance KnownShX '[] where knownShX = ZKX
instance (KnownNat n, KnownShX sh) => KnownShX (Just n : sh) where knownShX = SKnown natSing :!% knownShX
instance KnownShX sh => KnownShX (Nothing : sh) where knownShX = SUnknown () :!% knownShX

instance (KnownShX sh', KnownElt a) => KnownElt (Mixed sh' a) where
  memptyArray sh = M_Nest sh (memptyArray (X.shAppend sh (X.completeShXzeros (knownShX @sh'))))

  mvecsUnsafeNew sh example
    | X.shapeSize sh' == 0 = mvecsNewEmpty (Proxy @(Mixed sh' a))
    | otherwise = MV_Nest sh' <$> mvecsUnsafeNew (X.shAppend sh sh') (mindex example (X.zeroIxX (X.staticShapeFrom sh')))
    where
      sh' = mshape example

  mvecsNewEmpty _ = MV_Nest (X.completeShXzeros (knownShX @sh')) <$> mvecsNewEmpty (Proxy @a)


-- | Create an array given a size and a function that computes the element at a
-- given index.
--
-- __WARNING__: It is required that every @a@ returned by the argument to
-- 'mgenerate' has the same shape. For example, the following will throw a
-- runtime error:
--
-- > foo :: Mixed [Nothing] (Mixed [Nothing] Double)
-- > foo = mgenerate (10 :.: ZIR) $ \(i :.: ZIR) ->
-- >         mgenerate (i :.: ZIR) $ \(j :.: ZIR) ->
-- >           ...
--
-- because the size of the inner 'mgenerate' is not always the same (it depends
-- on @i@). Nested arrays in @ox-arrays@ are always stored fully flattened, so
-- the entire hierarchy (after distributing out tuples) must be a rectangular
-- array. The type of 'mgenerate' allows this requirement to be broken very
-- easily, hence the runtime check.
mgenerate :: forall sh a. KnownElt a => IShX sh -> (IIxX sh -> a) -> Mixed sh a
mgenerate sh f = case X.enumShape sh of
  [] -> memptyArray sh
  firstidx : restidxs ->
    let firstelem = f (X.zeroIxX' sh)
        shapetree = mshapeTree firstelem
    in if mshapeTreeEmpty (Proxy @a) shapetree
         then memptyArray sh
         else runST $ do
                vecs <- mvecsUnsafeNew sh firstelem
                mvecsWrite sh firstidx firstelem vecs
                -- TODO: This is likely fine if @a@ is big, but if @a@ is a
                -- scalar this array copying inefficient. Should improve this.
                forM_ restidxs $ \idx -> do
                  let val = f idx
                  when (not (mshapeTreeEq (Proxy @a) (mshapeTree val) shapetree)) $
                    error "Data.Array.Nested mgenerate: generated values do not have equal shapes"
                  mvecsWrite sh idx val vecs
                mvecsFreeze sh vecs

mtranspose :: forall is sh a. (X.Permutation is, X.Rank is <= X.Rank sh, Elt a)
           => HList SNat is -> Mixed sh a -> Mixed (X.PermutePrefix is sh) a
mtranspose perm arr =
  let ssh = X.staticShapeFrom (mshape arr)
      sshPP = X.ssxAppend (X.ssxPermute perm (X.ssxTakeLen perm ssh)) (X.ssxDropLen perm ssh)
  in mlift sshPP (\ssh' -> X.rerankTop ssh sshPP ssh' (X.transpose ssh perm)) arr

mappend :: forall n m sh a. Elt a
        => Mixed (n : sh) a -> Mixed (m : sh) a -> Mixed (X.AddMaybe n m : sh) a
mappend arr1 arr2 = mlift2 (snm :!% ssh) f arr1 arr2
  where
    sn :$% sh = mshape arr1
    sm :$% _ = mshape arr2
    ssh = X.staticShapeFrom sh
    snm :: SMayNat () SNat (X.AddMaybe n m)
    snm = case (sn, sm) of
            (SUnknown{}, _) -> SUnknown ()
            (SKnown{}, SUnknown{}) -> SUnknown ()
            (SKnown n, SKnown m) -> SKnown (X.plusSNat n m)

    f :: forall sh' b. Storable b
      => StaticShX sh' -> XArray (n : sh ++ sh') b -> XArray (m : sh ++ sh') b -> XArray (X.AddMaybe n m : sh ++ sh') b
    f ssh' = X.append (X.ssxAppend ssh ssh')

mfromVectorP :: forall sh a. Storable a => IShX sh -> VS.Vector a -> Mixed sh (Primitive a)
mfromVectorP sh v = M_Primitive sh (X.fromVector sh v)

mfromVector :: forall sh a. (Storable a, PrimElt a) => IShX sh -> VS.Vector a -> Mixed sh a
mfromVector sh v = fromPrimitive (mfromVectorP sh v)

mtoVectorP :: Storable a => Mixed sh (Primitive a) -> VS.Vector a
mtoVectorP (M_Primitive _ v) = X.toVector v

mtoVector :: (Storable a, PrimElt a) => Mixed sh a -> VS.Vector a
mtoVector arr = mtoVectorP (coerce toPrimitive arr)

mfromList :: Elt a => NonEmpty a -> Mixed '[Nothing] a
mfromList = mfromList1 . fmap mscalar

mtoList :: Elt a => Mixed '[n] a -> [a]
mtoList = map munScalar . mtoList1

munScalar :: Elt a => Mixed '[] a -> a
munScalar arr = mindex arr ZIX

mconstantP :: forall sh a. Storable a => IShX sh -> a -> Mixed sh (Primitive a)
mconstantP sh x = M_Primitive sh (X.constant sh x)

mconstant :: forall sh a. (Storable a, PrimElt a)
          => IShX sh -> a -> Mixed sh a
mconstant sh x = fromPrimitive (mconstantP sh x)

mslice :: Elt a => SNat i -> SNat n -> Mixed (Just (i + n + k) : sh) a -> Mixed (Just n : sh) a
mslice i n arr =
  let _ :$% sh = mshape arr
  in mlift (SKnown n :!% X.staticShapeFrom sh) (\_ -> X.slice i n) arr

msliceU :: Elt a => Int -> Int -> Mixed (Nothing : sh) a -> Mixed (Nothing : sh) a
msliceU i n arr = mlift (X.staticShapeFrom (mshape arr)) (\_ -> X.sliceU i n) arr

mrev1 :: Elt a => Mixed (n : sh) a -> Mixed (n : sh) a
mrev1 arr = mlift (X.staticShapeFrom (mshape arr)) (\_ -> X.rev1) arr

mreshape :: forall sh sh' a. Elt a => IShX sh' -> Mixed sh a -> Mixed sh' a
mreshape sh' arr =
  mlift (X.staticShapeFrom sh')
        (\sshIn -> X.reshapePartial (X.staticShapeFrom (mshape arr)) sshIn sh')
        arr

masXArrayPrimP :: Mixed sh (Primitive a) -> (IShX sh, XArray sh a)
masXArrayPrimP (M_Primitive sh arr) = (sh, arr)

masXArrayPrim :: PrimElt a => Mixed sh a -> (IShX sh, XArray sh a)
masXArrayPrim = masXArrayPrimP . toPrimitive

mfromXArrayPrimP :: StaticShX sh -> XArray sh a -> Mixed sh (Primitive a)
mfromXArrayPrimP ssh arr = M_Primitive (X.shape ssh arr) arr

mfromXArrayPrim :: PrimElt a => StaticShX sh -> XArray sh a -> Mixed sh a
mfromXArrayPrim = (fromPrimitive .) . mfromXArrayPrimP

mliftPrim :: (Storable a, PrimElt a)
          => (a -> a)
          -> Mixed sh a -> Mixed sh a
mliftPrim f (toPrimitive -> M_Primitive sh (X.XArray arr)) = fromPrimitive $ M_Primitive sh (X.XArray (S.mapA f arr))

mliftPrim2 :: (Storable a, PrimElt a)
           => (a -> a -> a)
           -> Mixed sh a -> Mixed sh a -> Mixed sh a
mliftPrim2 f (toPrimitive -> M_Primitive sh (X.XArray arr1)) (toPrimitive -> M_Primitive _ (X.XArray arr2)) =
  fromPrimitive $ M_Primitive sh (X.XArray (S.zipWithA f arr1 arr2))

instance (Storable a, Num a, PrimElt a) => Num (Mixed sh a) where
  (+) = mliftPrim2 (+)
  (-) = mliftPrim2 (-)
  (*) = mliftPrim2 (*)
  negate = mliftPrim negate
  abs = mliftPrim abs
  signum = mliftPrim signum
  fromInteger _ = error "Data.Array.Nested.fromIntegral: No singletons available, use explicit mconstant"


-- | A rank-typed array: the number of dimensions of the array (its /rank/) is
-- represented on the type level as a 'Nat'.
--
-- Valid elements of a ranked arrays are described by the 'Elt' type class.
-- Because 'Ranked' itself is also an instance of 'Elt', nested arrays are
-- supported (and are represented as a single, flattened, struct-of-arrays
-- array internally).
--
-- 'Ranked' is a newtype around a 'Mixed' of 'Nothing's.
type Ranked :: Nat -> Type -> Type
newtype Ranked n a = Ranked (Mixed (Replicate n Nothing) a)
deriving instance Show (Mixed (Replicate n Nothing) a) => Show (Ranked n a)

-- | A shape-typed array: the full shape of the array (the sizes of its
-- dimensions) is represented on the type level as a list of 'Nat's. Note that
-- these are "GHC.TypeLits" naturals, because we do not need induction over
-- them and we want very large arrays to be possible.
--
-- Like for 'Ranked', the valid elements are described by the 'Elt' type class,
-- and 'Shaped' itself is again an instance of 'Elt' as well.
--
-- 'Shaped' is a newtype around a 'Mixed' of 'Just's.
type Shaped :: [Nat] -> Type -> Type
newtype Shaped sh a = Shaped (Mixed (MapJust sh) a)
deriving instance Show (Mixed (MapJust sh) a) => Show (Shaped sh a)

-- just unwrap the newtype and defer to the general instance for nested arrays
newtype instance Mixed sh (Ranked n   a) = M_Ranked (Mixed sh (Mixed (Replicate n Nothing) a))
deriving instance Show (Mixed sh (Mixed (Replicate n Nothing) a)) => Show (Mixed sh (Ranked n a))
newtype instance Mixed sh (Shaped sh' a) = M_Shaped (Mixed sh (Mixed (MapJust sh'        ) a))
deriving instance Show (Mixed sh (Mixed (MapJust sh'        ) a)) => Show (Mixed sh (Shaped sh' a))

newtype instance MixedVecs s sh (Ranked n   a) = MV_Ranked (MixedVecs s sh (Mixed (Replicate n Nothing) a))
newtype instance MixedVecs s sh (Shaped sh' a) = MV_Shaped (MixedVecs s sh (Mixed (MapJust sh'        ) a))


-- 'Ranked' and 'Shaped' can already be used at the top level of an array nest;
-- these instances allow them to also be used as elements of arrays, thus
-- making them first-class in the API.
instance Elt a => Elt (Ranked n a) where
  mshape (M_Ranked arr) = mshape arr
  mindex (M_Ranked arr) i = Ranked (mindex arr i)

  mindexPartial :: forall sh sh'. Mixed (sh ++ sh') (Ranked n a) -> IIxX sh -> Mixed sh' (Ranked n a)
  mindexPartial (M_Ranked arr) i =
    coerce @(Mixed sh' (Mixed (Replicate n Nothing) a)) @(Mixed sh' (Ranked n a)) $
        mindexPartial arr i

  mscalar (Ranked x) = M_Ranked (M_Nest ZSX x)

  mfromList1 :: forall sh. NonEmpty (Mixed sh (Ranked n a)) -> Mixed (Nothing : sh) (Ranked n a)
  mfromList1 l = M_Ranked (mfromList1 (coerce l))

  mtoList1 :: forall m sh. Mixed (m : sh) (Ranked n a) -> [Mixed sh (Ranked n a)]
  mtoList1 (M_Ranked arr) =
    coerce @[Mixed sh (Mixed (Replicate n 'Nothing) a)] @[Mixed sh (Ranked n a)] (mtoList1 arr)

  mlift :: forall sh1 sh2.
           StaticShX sh2
        -> (forall sh' b. Storable b => StaticShX sh' -> XArray (sh1 ++ sh') b -> XArray (sh2 ++ sh') b)
        -> Mixed sh1 (Ranked n a) -> Mixed sh2 (Ranked n a)
  mlift ssh2 f (M_Ranked arr) =
    coerce @(Mixed sh2 (Mixed (Replicate n Nothing) a)) @(Mixed sh2 (Ranked n a)) $
      mlift ssh2 f arr

  mlift2 :: forall sh1 sh2 sh3.
            StaticShX sh3
         -> (forall sh' b. Storable b => StaticShX sh' -> XArray (sh1 ++ sh') b -> XArray (sh2 ++ sh') b -> XArray (sh3 ++ sh') b)
         -> Mixed sh1 (Ranked n a) -> Mixed sh2 (Ranked n a) -> Mixed sh3 (Ranked n a)
  mlift2 ssh3 f (M_Ranked arr1) (M_Ranked arr2) =
    coerce @(Mixed sh3 (Mixed (Replicate n Nothing) a)) @(Mixed sh3 (Ranked n a)) $
        mlift2 ssh3 f arr1 arr2

  mcast ssh1 sh2 psh' (M_Ranked arr) = M_Ranked (mcast ssh1 sh2 psh' arr)

  mshapeTree (Ranked arr) = first shCvtXR' (mshapeTree arr)

  mshapeTreeEq _ (sh1, t1) (sh2, t2) = sh1 == sh2 && mshapeTreeEq (Proxy @a) t1 t2

  mshapeTreeEmpty _ (sh, t) = shapeSizeR sh == 0 && mshapeTreeEmpty (Proxy @a) t

  mshowShapeTree _ (sh, t) = "(" ++ show sh ++ ", " ++ mshowShapeTree (Proxy @a) t ++ ")"

  mvecsWrite :: forall sh s. IShX sh -> IIxX sh -> Ranked n a -> MixedVecs s sh (Ranked n a) -> ST s ()
  mvecsWrite sh idx (Ranked arr) vecs =
    mvecsWrite sh idx arr
      (coerce @(MixedVecs s sh (Ranked n a)) @(MixedVecs s sh (Mixed (Replicate n Nothing) a))
         vecs)

  mvecsWritePartial :: forall sh sh' s.
                       IShX (sh ++ sh') -> IIxX sh -> Mixed sh' (Ranked n a)
                    -> MixedVecs s (sh ++ sh') (Ranked n a)
                    -> ST s ()
  mvecsWritePartial sh idx arr vecs =
    mvecsWritePartial sh idx
      (coerce @(Mixed sh' (Ranked n a))
              @(Mixed sh' (Mixed (Replicate n Nothing) a))
         arr)
      (coerce @(MixedVecs s (sh ++ sh') (Ranked n a))
              @(MixedVecs s (sh ++ sh') (Mixed (Replicate n Nothing) a))
         vecs)

  mvecsFreeze :: forall sh s. IShX sh -> MixedVecs s sh (Ranked n a) -> ST s (Mixed sh (Ranked n a))
  mvecsFreeze sh vecs =
    coerce @(Mixed sh (Mixed (Replicate n Nothing) a))
           @(Mixed sh (Ranked n a))
      <$> mvecsFreeze sh
            (coerce @(MixedVecs s sh (Ranked n a))
                    @(MixedVecs s sh (Mixed (Replicate n Nothing) a))
                    vecs)

instance (KnownNat n, KnownElt a) => KnownElt (Ranked n a) where
  memptyArray :: forall sh. IShX sh -> Mixed sh (Ranked n a)
  memptyArray i
    | Dict <- lemKnownReplicate (SNat @n)
    = coerce @(Mixed sh (Mixed (Replicate n Nothing) a)) @(Mixed sh (Ranked n a)) $
        memptyArray i

  mvecsUnsafeNew idx (Ranked arr)
    | Dict <- lemKnownReplicate (SNat @n)
    = MV_Ranked <$> mvecsUnsafeNew idx arr

  mvecsNewEmpty _
    | Dict <- lemKnownReplicate (SNat @n)
    = MV_Ranked <$> mvecsNewEmpty (Proxy @(Mixed (Replicate n Nothing) a))

-- sshapeKnown :: ShS sh -> Dict KnownShape sh
-- sshapeKnown ZSS = Dict
-- sshapeKnown (GHC_SNat :$$ sh) | Dict <- sshapeKnown sh = Dict

lemCommMapJustApp :: forall sh1 sh2. ShS sh1 -> Proxy sh2
                  -> MapJust (sh1 ++ sh2) :~: MapJust sh1 ++ MapJust sh2
lemCommMapJustApp ZSS _ = Refl
lemCommMapJustApp (_ :$$ sh) p | Refl <- lemCommMapJustApp sh p = Refl

instance Elt a => Elt (Shaped sh a) where
  mshape (M_Shaped arr) = mshape arr
  mindex (M_Shaped arr) i = Shaped (mindex arr i)

  mindexPartial :: forall sh1 sh2. Mixed (sh1 ++ sh2) (Shaped sh a) -> IIxX sh1 -> Mixed sh2 (Shaped sh a)
  mindexPartial (M_Shaped arr) i =
    coerce @(Mixed sh2 (Mixed (MapJust sh) a)) @(Mixed sh2 (Shaped sh a)) $
      mindexPartial arr i

  mscalar (Shaped x) = M_Shaped (M_Nest ZSX x)

  mfromList1 :: forall sh'. NonEmpty (Mixed sh' (Shaped sh a)) -> Mixed (Nothing : sh') (Shaped sh a)
  mfromList1 l = M_Shaped (mfromList1 (coerce l))

  mtoList1 :: forall n sh'. Mixed (n : sh') (Shaped sh a) -> [Mixed sh' (Shaped sh a)]
  mtoList1 (M_Shaped arr)
    = coerce @[Mixed sh' (Mixed (MapJust sh) a)] @[Mixed sh' (Shaped sh a)] (mtoList1 arr)

  mlift :: forall sh1 sh2.
           StaticShX sh2
        -> (forall sh' b. Storable b => StaticShX sh' -> XArray (sh1 ++ sh') b -> XArray (sh2 ++ sh') b)
        -> Mixed sh1 (Shaped sh a) -> Mixed sh2 (Shaped sh a)
  mlift ssh2 f (M_Shaped arr) =
    coerce @(Mixed sh2 (Mixed (MapJust sh) a)) @(Mixed sh2 (Shaped sh a)) $
      mlift ssh2 f arr

  mlift2 :: forall sh1 sh2 sh3.
            StaticShX sh3
         -> (forall sh' b. Storable b => StaticShX sh' -> XArray (sh1 ++ sh') b -> XArray (sh2 ++ sh') b -> XArray (sh3 ++ sh') b)
         -> Mixed sh1 (Shaped sh a) -> Mixed sh2 (Shaped sh a) -> Mixed sh3 (Shaped sh a)
  mlift2 ssh3 f (M_Shaped arr1) (M_Shaped arr2) =
    coerce @(Mixed sh3 (Mixed (MapJust sh) a)) @(Mixed sh3 (Shaped sh a)) $
      mlift2 ssh3 f arr1 arr2

  mcast ssh1 sh2 psh' (M_Shaped arr) = M_Shaped (mcast ssh1 sh2 psh' arr)

  mshapeTree (Shaped arr) = first shCvtXS' (mshapeTree arr)

  mshapeTreeEq _ (sh1, t1) (sh2, t2) = sh1 == sh2 && mshapeTreeEq (Proxy @a) t1 t2

  mshapeTreeEmpty _ (sh, t) = shapeSizeS sh == 0 && mshapeTreeEmpty (Proxy @a) t

  mshowShapeTree _ (sh, t) = "(" ++ show sh ++ ", " ++ mshowShapeTree (Proxy @a) t ++ ")"

  mvecsWrite :: forall sh' s. IShX sh' -> IIxX sh' -> Shaped sh a -> MixedVecs s sh' (Shaped sh a) -> ST s ()
  mvecsWrite sh idx (Shaped arr) vecs =
    mvecsWrite sh idx arr
      (coerce @(MixedVecs s sh' (Shaped sh a)) @(MixedVecs s sh' (Mixed (MapJust sh) a))
         vecs)

  mvecsWritePartial :: forall sh1 sh2 s.
                       IShX (sh1 ++ sh2) -> IIxX sh1 -> Mixed sh2 (Shaped sh a)
                    -> MixedVecs s (sh1 ++ sh2) (Shaped sh a)
                    -> ST s ()
  mvecsWritePartial sh idx arr vecs =
    mvecsWritePartial sh idx
      (coerce @(Mixed sh2 (Shaped sh a))
              @(Mixed sh2 (Mixed (MapJust sh) a))
         arr)
      (coerce @(MixedVecs s (sh1 ++ sh2) (Shaped sh a))
              @(MixedVecs s (sh1 ++ sh2) (Mixed (MapJust sh) a))
         vecs)

  mvecsFreeze :: forall sh' s. IShX sh' -> MixedVecs s sh' (Shaped sh a) -> ST s (Mixed sh' (Shaped sh a))
  mvecsFreeze sh vecs =
    coerce @(Mixed sh' (Mixed (MapJust sh) a))
           @(Mixed sh' (Shaped sh a))
      <$> mvecsFreeze sh
            (coerce @(MixedVecs s sh' (Shaped sh a))
                    @(MixedVecs s sh' (Mixed (MapJust sh) a))
                    vecs)

-- | Evidence for the static part of a shape. This pops up only when you are
-- polymorphic in the element type of an array.
type KnownShS :: [Nat] -> Constraint
class KnownShS sh where knownShS :: ShS sh
instance KnownShS '[] where knownShS = ZSS
instance (KnownNat n, KnownShS sh) => KnownShS (n : sh) where knownShS = natSing :$$ knownShS

lemKnownMapJust :: forall sh. KnownShS sh => Proxy sh -> Dict KnownShX (MapJust sh)
lemKnownMapJust _ = lemKnownShX (go (knownShS @sh))
  where
    go :: ShS sh' -> StaticShX (MapJust sh')
    go ZSS = ZKX
    go (n :$$ sh) = SKnown n :!% go sh

instance (KnownShS sh, KnownElt a) => KnownElt (Shaped sh a) where
  memptyArray :: forall sh'. IShX sh' -> Mixed sh' (Shaped sh a)
  memptyArray i
    | Dict <- lemKnownMapJust (Proxy @sh)
    = coerce @(Mixed sh' (Mixed (MapJust sh) a)) @(Mixed sh' (Shaped sh a)) $
        memptyArray i

  mvecsUnsafeNew idx (Shaped arr)
    | Dict <- lemKnownMapJust (Proxy @sh)
    = MV_Shaped <$> mvecsUnsafeNew idx arr

  mvecsNewEmpty _
    | Dict <- lemKnownMapJust (Proxy @sh)
    = MV_Shaped <$> mvecsNewEmpty (Proxy @(Mixed (MapJust sh) a))


-- ====== API OF RANKED ARRAYS ====== --

arithPromoteRanked :: forall n a. PrimElt a
                   => (forall sh. Mixed sh a -> Mixed sh a)
                   -> Ranked n a -> Ranked n a
arithPromoteRanked = coerce

arithPromoteRanked2 :: forall n a. PrimElt a
                    => (forall sh. Mixed sh a -> Mixed sh a -> Mixed sh a)
                    -> Ranked n a -> Ranked n a -> Ranked n a
arithPromoteRanked2 = coerce

instance (Storable a, Num a, PrimElt a) => Num (Ranked n a) where
  (+) = arithPromoteRanked2 (+)
  (-) = arithPromoteRanked2 (-)
  (*) = arithPromoteRanked2 (*)
  negate = arithPromoteRanked negate
  abs = arithPromoteRanked abs
  signum = arithPromoteRanked signum
  fromInteger _ = error "Data.Array.Nested.fromIntegral: No singletons available, use explicit mconstant"

zeroIxR :: SNat n -> IIxR n
zeroIxR SZ = ZIR
zeroIxR (SS n) = 0 :.: zeroIxR n

ixCvtXR :: IIxX sh -> IIxR (X.Rank sh)
ixCvtXR ZIX = ZIR
ixCvtXR (n :.% idx) = n :.: ixCvtXR idx

shCvtXR' :: forall n. IShX (Replicate n Nothing) -> IShR n
shCvtXR' ZSX =
  castWith (subst2 (unsafeCoerce Refl :: 0 :~: n))
    ZSR
shCvtXR' (n :$% (idx :: IShX sh))
  | Refl <- lemReplicateSucc @(Nothing @Nat) @(n - 1) =
  castWith (subst2 (lem1 @sh Refl))
    (X.fromSMayNat' n :$: shCvtXR' (castWith (subst2 (lem2 Refl)) idx))
  where
    lem1 :: forall sh' n' k.
            k : sh' :~: Replicate n' Nothing
         -> Rank sh' + 1 :~: n'
    lem1 Refl = unsafeCoerce Refl

    lem2 :: k : sh :~: Replicate n Nothing
         -> sh :~: Replicate (Rank sh) Nothing
    lem2 Refl = unsafeCoerce Refl

ixCvtRX :: IIxR n -> IIxX (Replicate n Nothing)
ixCvtRX ZIR = ZIX
ixCvtRX (n :.: (idx :: IxR m Int)) =
  castWith (subst2 @IxX @Int (X.lemReplicateSucc @(Nothing @Nat) @m))
    (n :.% ixCvtRX idx)

shCvtRX :: IShR n -> IShX (Replicate n Nothing)
shCvtRX ZSR = ZSX
shCvtRX (n :$: (idx :: ShR m Int)) =
  castWith (subst2 @ShX @Int (X.lemReplicateSucc @(Nothing @Nat) @m))
    (SUnknown n :$% shCvtRX idx)

shapeSizeR :: IShR n -> Int
shapeSizeR ZSR = 1
shapeSizeR (n :$: sh) = n * shapeSizeR sh


rshape :: forall n a. Elt a => Ranked n a -> IShR n
rshape (Ranked arr) = shCvtXR' (mshape arr)

rindex :: Elt a => Ranked n a -> IIxR n -> a
rindex (Ranked arr) idx = mindex arr (ixCvtRX idx)

snatFromListR :: ListR n i -> SNat n
snatFromListR ZR = SNat
snatFromListR (_ ::: (l :: ListR n i)) | SNat <- snatFromListR l, Dict <- knownNatSucc @n = SNat

snatFromIxR :: IxR n i -> SNat n
snatFromIxR (IxR sh) = snatFromListR sh

snatFromShR :: ShR n i -> SNat n
snatFromShR (ShR sh) = snatFromListR sh

rindexPartial :: forall n m a. Elt a => Ranked (n + m) a -> IIxR n -> Ranked m a
rindexPartial (Ranked arr) idx =
  Ranked (mindexPartial @a @(Replicate n Nothing) @(Replicate m Nothing)
            (castWith (subst2 (lemReplicatePlusApp (snatFromIxR idx) (Proxy @m) (Proxy @Nothing))) arr)
            (ixCvtRX idx))

-- | __WARNING__: All values returned from the function must have equal shape.
-- See the documentation of 'mgenerate' for more details.
rgenerate :: forall n a. KnownElt a => IShR n -> (IIxR n -> a) -> Ranked n a
rgenerate sh f
  | sn@SNat <- snatFromShR sh
  , Dict <- lemKnownReplicate sn
  , Refl <- lemRankReplicate sn
  = Ranked (mgenerate (shCvtRX sh) (f . ixCvtXR))

-- | See the documentation of 'mlift'.
rlift :: forall n1 n2 a. Elt a
      => SNat n2
      -> (forall sh' b. Storable b => StaticShX sh' -> XArray (Replicate n1 Nothing ++ sh') b -> XArray (Replicate n2 Nothing ++ sh') b)
      -> Ranked n1 a -> Ranked n2 a
rlift sn2 f (Ranked arr) = Ranked (mlift (ssxFromSNat sn2) f arr)

rsumOuter1P :: forall n a.
               (Storable a, Num a)
            => Ranked (n + 1) (Primitive a) -> Ranked n (Primitive a)
rsumOuter1P (Ranked (M_Primitive sh arr))
  | Refl <- X.lemReplicateSucc @(Nothing @Nat) @n
  , _ :$% shT <- sh
  = Ranked (M_Primitive shT (X.sumOuter (SUnknown () :!% ZKX) (X.staticShapeFrom shT) arr))

rsumOuter1 :: forall n a. (Storable a, Num a, PrimElt a)
           => Ranked (n + 1) a -> Ranked n a
rsumOuter1 = coerce fromPrimitive . rsumOuter1P @n @a . coerce toPrimitive

rtranspose :: forall n a. Elt a => [Int] -> Ranked n a -> Ranked n a
rtranspose perm arr
  | sn@SNat <- snatFromShR (rshape arr)
  , Dict <- lemKnownReplicate sn
  , length perm <= fromIntegral (natVal (Proxy @n))
  = rlift sn
          (\ssh' -> X.transposeUntyped (natSing @n) ssh' perm)
          arr
  | otherwise
  = error "Data.Array.Nested.rtranspose: Permutation longer than rank of array"

rappend :: forall n a. Elt a
        => Ranked (n + 1) a -> Ranked (n + 1) a -> Ranked (n + 1) a
rappend arr1 arr2
  | sn@SNat <- snatFromShR (rshape arr1)
  , Dict <- lemKnownReplicate sn
  , Refl <- X.lemReplicateSucc @(Nothing @Nat) @n
  = coerce (mappend @Nothing @Nothing @(Replicate n Nothing))
      arr1 arr2

rscalar :: Elt a => a -> Ranked 0 a
rscalar x = Ranked (mscalar x)

rfromVectorP :: forall n a. Storable a => IShR n -> VS.Vector a -> Ranked n (Primitive a)
rfromVectorP sh v
  | Dict <- lemKnownReplicate (snatFromShR sh)
  = Ranked (mfromVectorP (shCvtRX sh) v)

rfromVector :: forall n a. (Storable a, PrimElt a) => IShR n -> VS.Vector a -> Ranked n a
rfromVector sh v = coerce fromPrimitive (rfromVectorP sh v)

rtoVectorP :: Storable a => Ranked n (Primitive a) -> VS.Vector a
rtoVectorP = coerce mtoVectorP

rtoVector :: (Storable a, PrimElt a) => Ranked n a -> VS.Vector a
rtoVector = coerce mtoVector

rfromList1 :: forall n a. Elt a => NonEmpty (Ranked n a) -> Ranked (n + 1) a
rfromList1 l
  | Refl <- X.lemReplicateSucc @(Nothing @Nat) @n
  = Ranked (mfromList1 (coerce l :: NonEmpty (Mixed (Replicate n Nothing) a)))

rfromList :: Elt a => NonEmpty a -> Ranked 1 a
rfromList l = Ranked (mfromList l)

rtoList :: forall n a. Elt a => Ranked (n + 1) a -> [Ranked n a]
rtoList (Ranked arr)
  | Refl <- X.lemReplicateSucc @(Nothing @Nat) @n
  = coerce (mtoList1 @a @Nothing @(Replicate n Nothing) arr)

rtoList1 :: Elt a => Ranked 1 a -> [a]
rtoList1 = map runScalar . rtoList

runScalar :: Elt a => Ranked 0 a -> a
runScalar arr = rindex arr ZIR

rconstantP :: forall n a. Storable a => IShR n -> a -> Ranked n (Primitive a)
rconstantP sh x
  | Dict <- lemKnownReplicate (snatFromShR sh)
  = Ranked (mconstantP (shCvtRX sh) x)

rconstant :: forall n a. (Storable a, PrimElt a)
          => IShR n -> a -> Ranked n a
rconstant sh x = coerce fromPrimitive (rconstantP sh x)

rslice :: forall n a. Elt a => Int -> Int -> Ranked (n + 1) a -> Ranked (n + 1) a
rslice i n arr
  | Refl <- X.lemReplicateSucc @(Nothing @Nat) @n
  = rlift (snatFromShR (rshape arr))
          (\_ -> X.sliceU i n)
          arr

rrev1 :: forall n a. Elt a => Ranked (n + 1) a -> Ranked (n + 1) a
rrev1 arr =
  rlift (snatFromShR (rshape arr))
        (\(_ :: StaticShX sh') ->
          case X.lemReplicateSucc @(Nothing @Nat) @n of
            Refl -> X.rev1 @Nothing @(Replicate n Nothing ++ sh'))
        arr

rreshape :: forall n n' a. Elt a
         => IShR n' -> Ranked n a -> Ranked n' a
rreshape sh' rarr@(Ranked arr)
  | Dict <- lemKnownReplicate (snatFromShR (rshape rarr))
  , Dict <- lemKnownReplicate (snatFromShR sh')
  = Ranked (mreshape (shCvtRX sh') arr)

rasXArrayPrimP :: Ranked n (Primitive a) -> (IShR n, XArray (Replicate n Nothing) a)
rasXArrayPrimP (Ranked arr) = first shCvtXR' (masXArrayPrimP arr)

rasXArrayPrim :: PrimElt a => Ranked n a -> (IShR n, XArray (Replicate n Nothing) a)
rasXArrayPrim (Ranked arr) = first shCvtXR' (masXArrayPrim arr)

rfromXArrayPrimP :: SNat n -> XArray (Replicate n Nothing) a -> Ranked n (Primitive a)
rfromXArrayPrimP sn arr = Ranked (mfromXArrayPrimP (X.staticShapeFrom (X.shape (ssxFromSNat sn) arr)) arr)

rfromXArrayPrim :: PrimElt a => SNat n -> XArray (Replicate n Nothing) a -> Ranked n a
rfromXArrayPrim sn arr = Ranked (mfromXArrayPrim (X.staticShapeFrom (X.shape (ssxFromSNat sn) arr)) arr)


-- ====== API OF SHAPED ARRAYS ====== --

arithPromoteShaped :: forall sh a. PrimElt a
                   => (forall shx. Mixed shx a -> Mixed shx a)
                   -> Shaped sh a -> Shaped sh a
arithPromoteShaped = coerce

arithPromoteShaped2 :: forall sh a. PrimElt a
                    => (forall shx. Mixed shx a -> Mixed shx a -> Mixed shx a)
                    -> Shaped sh a -> Shaped sh a -> Shaped sh a
arithPromoteShaped2 = coerce

instance (Storable a, Num a, PrimElt a) => Num (Shaped sh a) where
  (+) = arithPromoteShaped2 (+)
  (-) = arithPromoteShaped2 (-)
  (*) = arithPromoteShaped2 (*)
  negate = arithPromoteShaped negate
  abs = arithPromoteShaped abs
  signum = arithPromoteShaped signum
  fromInteger _ = error "Data.Array.Nested.fromIntegral: No singletons available, use explicit mconstant"

zeroIxS :: ShS sh -> IIxS sh
zeroIxS ZSS = ZIS
zeroIxS (_ :$$ sh) = 0 :.$ zeroIxS sh

ixCvtXS :: ShS sh -> IIxX (MapJust sh) -> IIxS sh
ixCvtXS ZSS ZIX = ZIS
ixCvtXS (_ :$$ sh) (n :.% idx) = n :.$ ixCvtXS sh idx

type family Tail l where
  Tail (_ : xs) = xs

shCvtXS' :: forall sh. IShX (MapJust sh) -> ShS sh
shCvtXS' ZSX = castWith (subst1 (unsafeCoerce Refl :: '[] :~: sh)) ZSS
shCvtXS' (SKnown n :$% (idx :: IShX mjshT)) =
  castWith (subst1 (lem Refl)) $
    n :$$ shCvtXS' @(Tail sh) (castWith (subst2 (unsafeCoerce Refl :: mjshT :~: MapJust (Tail sh)))
                                 idx)
  where
    lem :: forall sh1 sh' n.
           Just n : sh1 :~: MapJust sh'
        -> n : Tail sh' :~: sh'
    lem Refl = unsafeCoerce Refl
shCvtXS' (SUnknown _ :$% _) = error "impossible"

ixCvtSX :: IIxS sh -> IIxX (MapJust sh)
ixCvtSX ZIS = ZIX
ixCvtSX (n :.$ sh) = n :.% ixCvtSX sh

shCvtSX :: ShS sh -> IShX (MapJust sh)
shCvtSX ZSS = ZSX
shCvtSX (n :$$ sh) = SKnown n :$% shCvtSX sh

shapeSizeS :: ShS sh -> Int
shapeSizeS ZSS = 1
shapeSizeS (n :$$ sh) = X.fromSNat' n * shapeSizeS sh


sshape :: forall sh a. Elt a => Shaped sh a -> ShS sh
sshape (Shaped arr) = shCvtXS' (mshape arr)

sindex :: Elt a => Shaped sh a -> IIxS sh -> a
sindex (Shaped arr) idx = mindex arr (ixCvtSX idx)

shsTakeIx :: Proxy sh' -> ShS (sh ++ sh') -> IIxS sh -> ShS sh
shsTakeIx _ _ ZIS = ZSS
shsTakeIx p sh (_ :.$ idx) = case sh of n :$$ sh' -> n :$$ shsTakeIx p sh' idx

sindexPartial :: forall sh1 sh2 a. Elt a => Shaped (sh1 ++ sh2) a -> IIxS sh1 -> Shaped sh2 a
sindexPartial sarr@(Shaped arr) idx =
  Shaped (mindexPartial @a @(MapJust sh1) @(MapJust sh2)
            (castWith (subst2 (lemCommMapJustApp (shsTakeIx (Proxy @sh2) (sshape sarr) idx) (Proxy @sh2))) arr)
            (ixCvtSX 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 sh f = Shaped (mgenerate (shCvtSX sh) (f . ixCvtXS sh))

{-
-- | See the documentation of 'mlift'.
slift :: forall sh1 sh2 a. (KnownShape sh2, Elt a)
      => (forall sh' b. (KnownShapeX sh', Storable b) => Proxy sh' -> XArray (MapJust sh1 ++ sh') b -> XArray (MapJust sh2 ++ sh') b)
      -> Shaped sh1 a -> Shaped sh2 a
slift f (Shaped arr)
  | Dict <- lemKnownMapJust (Proxy @sh2)
  = Shaped (mlift f arr)

ssumOuter1P :: forall sh n a.
               (Storable a, Num a, KnownNat n, KnownShape sh)
            => Shaped (n : sh) (Primitive a) -> Shaped sh (Primitive a)
ssumOuter1P (Shaped arr)
  | Dict <- lemKnownMapJust (Proxy @sh)
  = Shaped
    . coerce @(XArray (MapJust sh) a) @(Mixed (MapJust sh) (Primitive a))
    . X.sumOuter (natSing @n :!$@ ZKSX) (knownShapeX @(MapJust sh))
    . coerce @(Mixed (Just n : MapJust sh) (Primitive a)) @(XArray (Just n : MapJust sh) a)
    $ arr

ssumOuter1 :: forall sh n a.
              (Storable a, Num a, PrimElt a, KnownNat n, KnownShape sh)
           => Shaped (n : sh) a -> Shaped sh a
ssumOuter1 = coerce fromPrimitive . ssumOuter1P @sh @n @a . coerce toPrimitive

lemCommMapJustTakeLen :: HList SNat is -> ShS sh -> X.TakeLen is (MapJust sh) :~: MapJust (X.TakeLen is sh)
lemCommMapJustTakeLen HNil _ = Refl
lemCommMapJustTakeLen (_ `HCons` is) (_ :$$ sh) | Refl <- lemCommMapJustTakeLen is sh = Refl
lemCommMapJustTakeLen (_ `HCons` _) ZSS = error "TakeLen of empty"

lemCommMapJustDropLen :: HList SNat is -> ShS sh -> X.DropLen is (MapJust sh) :~: MapJust (X.DropLen is sh)
lemCommMapJustDropLen HNil _ = Refl
lemCommMapJustDropLen (_ `HCons` is) (_ :$$ sh) | Refl <- lemCommMapJustDropLen is sh = Refl
lemCommMapJustDropLen (_ `HCons` _) ZSS = error "DropLen of empty"

lemCommMapJustIndex :: SNat i -> ShS sh -> X.Index i (MapJust sh) :~: Just (X.Index i sh)
lemCommMapJustIndex SZ (_ :$$ _) = Refl
lemCommMapJustIndex (SS (i :: SNat i')) ((_ :: SNat n) :$$ (sh :: ShS sh'))
  | Refl <- lemCommMapJustIndex i sh
  , Refl <- X.lemIndexSucc (Proxy @i') (Proxy @(Just n)) (Proxy @(MapJust sh'))
  , Refl <- X.lemIndexSucc (Proxy @i') (Proxy @n) (Proxy @sh')
  = Refl
lemCommMapJustIndex _ ZSS = error "Index of empty"

lemCommMapJustPermute :: HList SNat is -> ShS sh -> X.Permute is (MapJust sh) :~: MapJust (X.Permute is sh)
lemCommMapJustPermute HNil _ = Refl
lemCommMapJustPermute (i `HCons` is) sh
  | Refl <- lemCommMapJustPermute is sh
  , Refl <- lemCommMapJustIndex i sh
  = Refl

shTakeLen :: HList SNat is -> ShS sh -> ShS (X.TakeLen is sh)
shTakeLen HNil _ = ZSS
shTakeLen (_ `HCons` is) (n :$$ sh) = n :$$ shTakeLen is sh
shTakeLen (_ `HCons` _) ZSS = error "Permutation longer than shape"

shPermute :: HList SNat is -> ShS sh -> ShS (X.Permute is sh)
shPermute HNil _ = ZSS
shPermute (i `HCons` (is :: HList SNat is')) (sh :: ShS sh) = shIndex (Proxy @is') (Proxy @sh) i sh (shPermute is sh)

shIndex :: Proxy is -> Proxy shT -> SNat i -> ShS sh -> ShS (X.Permute is shT) -> ShS (X.Index i sh : X.Permute is shT)
shIndex _ _ SZ (n :$$ _) rest = n :$$ rest
shIndex p pT (SS (i :: SNat i')) ((_ :: SNat n) :$$ (sh :: ShS sh')) rest
  | Refl <- X.lemIndexSucc (Proxy @i') (Proxy @n) (Proxy @sh')
  = shIndex p pT i sh rest
shIndex _ _ _ ZSS _ = error "Index into empty shape"

stranspose :: forall is sh a. (X.Permutation is, X.Rank is <= X.Rank sh, KnownShape sh, Elt a) => HList SNat is -> Shaped sh a -> Shaped (X.PermutePrefix is sh) a
stranspose perm (Shaped arr)
  | Dict <- lemKnownMapJust (Proxy @sh)
  , Refl <- lemRankMapJust (Proxy @sh)
  , Refl <- lemCommMapJustTakeLen perm (knownShape @sh)
  , Refl <- lemCommMapJustDropLen perm (knownShape @sh)
  , Refl <- lemCommMapJustPermute perm (shTakeLen perm (knownShape @sh))
  , Refl <- lemCommMapJustApp (shPermute perm (shTakeLen perm (knownShape @sh))) (Proxy @(X.DropLen is sh))
  = Shaped (mtranspose perm arr)

sappend :: forall n m sh a. (KnownNat n, KnownNat m, KnownShape sh, Elt a)
        => Shaped (n : sh) a -> Shaped (m : sh) a -> Shaped (n + m : sh) a
sappend | Dict <- lemKnownMapJust (Proxy @sh) = coerce mappend

sscalar :: Elt a => a -> Shaped '[] a
sscalar x = Shaped (mscalar x)

sfromVectorP :: forall sh a. (KnownShape sh, Storable a) => VS.Vector a -> Shaped sh (Primitive a)
sfromVectorP v
  | Dict <- lemKnownMapJust (Proxy @sh)
  = Shaped (mfromVectorP (shCvtSX (knownShape @sh)) v)

sfromVector :: forall sh a. (KnownShape sh, Storable a, PrimElt a) => VS.Vector a -> Shaped sh a
sfromVector v = coerce fromPrimitive (sfromVectorP @sh @a v)

stoVectorP :: Storable a => Shaped sh (Primitive a) -> VS.Vector a
stoVectorP = coerce mtoVectorP

stoVector :: (Storable a, PrimElt a) => Shaped sh a -> VS.Vector a
stoVector = coerce mtoVector

sfromList1 :: forall n sh a. (KnownNat n, KnownShape sh, Elt a)
           => NonEmpty (Shaped sh a) -> Shaped (n : sh) a
sfromList1 l
  | Dict <- lemKnownMapJust (Proxy @sh)
  = Shaped (mfromList1 (coerce l))

sfromList :: (KnownNat n, Elt a) => NonEmpty a -> Shaped '[n] a
sfromList = Shaped . mfromList1 . fmap mscalar

stoList :: Elt a => Shaped (n : sh) a -> [Shaped sh a]
stoList (Shaped arr) = coerce (mtoList1 arr)

stoList1 :: Elt a => Shaped '[n] a -> [a]
stoList1 = map sunScalar . stoList

sunScalar :: Elt a => Shaped '[] a -> a
sunScalar arr = sindex arr ZIS

sconstantP :: forall sh a. (KnownShape sh, Storable a) => a -> Shaped sh (Primitive a)
sconstantP x
  | Dict <- lemKnownMapJust (Proxy @sh)
  = Shaped (mconstantP (shCvtSX (knownShape @sh)) x)

sconstant :: forall sh a. (KnownShape sh, Storable a, PrimElt a)
          => a -> Shaped sh a
sconstant x = coerce fromPrimitive (sconstantP @sh x)

sslice :: (KnownShape sh, Elt a) => SNat i -> SNat n -> Shaped (i + n + k : sh) a -> Shaped (n : sh) a
sslice i n@SNat = slift $ \_ -> X.slice i n

srev1 :: (KnownNat n, KnownShape sh, Elt a) => Shaped (n : sh) a -> Shaped (n : sh) a
srev1 = slift $ \_ -> X.rev1

sreshape :: forall sh sh' a. (KnownShape sh, KnownShape sh', Elt a)
         => ShS sh' -> Shaped sh a -> Shaped sh' a
sreshape sh' (Shaped arr)
  | Dict <- lemKnownMapJust (Proxy @sh)
  , Dict <- lemKnownMapJust (Proxy @sh')
  = Shaped (mreshape (shCvtSX sh') arr)

sasXArrayPrimP :: Shaped sh (Primitive a) -> XArray (MapJust sh) a
sasXArrayPrimP (Shaped arr) = masXArrayPrimP arr

sasXArrayPrim :: PrimElt a => Shaped sh a -> XArray (MapJust sh) a
sasXArrayPrim (Shaped arr) = masXArrayPrim arr

sfromXArrayPrimP :: XArray (MapJust sh) a -> Shaped sh (Primitive a)
sfromXArrayPrimP = Shaped . mfromXArrayPrimP

sfromXArrayPrim :: PrimElt a => XArray (MapJust sh) a -> Shaped sh a
sfromXArrayPrim = Shaped . mfromXArrayPrim
-}