1 files changed, 31 insertions, 21 deletions
diff --git a/src/Data/Array/Nested/Shaped/Base.hs b/src/Data/Array/Nested/Shaped/Base.hs
index 98f1241..4b119c4 100644
--- a/src/Data/Array/Nested/Shaped/Base.hs
+++ b/src/Data/Array/Nested/Shaped/Base.hs
@@ -26,7 +26,6 @@ import Data.Coerce (coerce)
 import Data.Kind (Type)
 import Data.List.NonEmpty (NonEmpty)
 import Data.Proxy
-import Data.Type.Equality
 import Foreign.Storable (Storable)
 import GHC.Float qualified (expm1, log1mexp, log1p, log1pexp)
 import GHC.Generics (Generic)
@@ -80,9 +79,12 @@ deriving instance Eq (Mixed sh (Mixed (MapJust sh') a)) => Eq (Mixed sh (Shaped
 newtype instance MixedVecs s sh (Shaped sh' a) = MV_Shaped (MixedVecs s sh (Mixed (MapJust sh') a))
 
 instance Elt a => Elt (Shaped sh a) where
+  {-# INLINE mshape #-}
   mshape (M_Shaped arr) = mshape arr
+  {-# INLINE mindex #-}
   mindex (M_Shaped arr) i = Shaped (mindex arr i)
 
+  {-# INLINE mindexPartial #-}
   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)) $
@@ -97,6 +99,7 @@ instance Elt a => Elt (Shaped sh a) where
   mtoListOuter (M_Shaped arr)
     = coerce @[Mixed sh' (Mixed (MapJust sh) a)] @[Mixed sh' (Shaped sh a)] (mtoListOuter arr)
 
+  {-# INLINE mlift #-}
   mlift :: forall sh1 sh2.
            StaticShX sh2
         -> (forall sh' b. Storable b => StaticShX sh' -> XArray (sh1 ++ sh') b -> XArray (sh2 ++ sh') b)
@@ -105,6 +108,7 @@ instance Elt a => Elt (Shaped sh a) where
     coerce @(Mixed sh2 (Mixed (MapJust sh) a)) @(Mixed sh2 (Shaped sh a)) $
       mlift ssh2 f arr
 
+  {-# INLINE mlift2 #-}
   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)
@@ -113,6 +117,7 @@ instance Elt a => Elt (Shaped sh a) where
     coerce @(Mixed sh3 (Mixed (MapJust sh) a)) @(Mixed sh3 (Shaped sh a)) $
       mlift2 ssh3 f arr1 arr2
 
+  {-# INLINE mliftL #-}
   mliftL :: forall sh1 sh2.
             StaticShX sh2
          -> (forall sh' b. Storable b => StaticShX sh' -> NonEmpty (XArray (sh1 ++ sh') b) -> NonEmpty (XArray (sh2 ++ sh') b))
@@ -132,7 +137,7 @@ instance Elt a => Elt (Shaped sh a) where
 
   type ShapeTree (Shaped sh a) = (ShS sh, ShapeTree a)
 
-  mshapeTree (Shaped arr) = first shsFromShX (mshapeTree arr)
+  mshapeTree (Shaped arr) = first coerce (mshapeTree arr)
 
   mshapeTreeEq _ (sh1, t1) (sh2, t2) = sh1 == sh2 && mshapeTreeEq (Proxy @a) t1 t2
 
@@ -142,18 +147,19 @@ instance Elt a => Elt (Shaped sh a) where
 
   marrayStrides (M_Shaped arr) = marrayStrides arr
 
-  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
+  mvecsWriteLinear :: forall sh' s. Int -> Shaped sh a -> MixedVecs s sh' (Shaped sh a) -> ST s ()
+  mvecsWriteLinear idx (Shaped arr) vecs =
+    mvecsWriteLinear 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
+  mvecsWritePartialLinear
+    :: forall sh1 sh2 s.
+       Proxy sh1 -> Int -> Mixed sh2 (Shaped sh a)
+    -> MixedVecs s (sh1 ++ sh2) (Shaped sh a)
+    -> ST s ()
+  mvecsWritePartialLinear proxy idx arr vecs =
+    mvecsWritePartialLinear proxy idx
       (coerce @(Mixed sh2 (Shaped sh a))
               @(Mixed sh2 (Mixed (MapJust sh) a))
          arr)
@@ -169,6 +175,14 @@ instance Elt a => Elt (Shaped sh a) where
             (coerce @(MixedVecs s sh' (Shaped sh a))
                     @(MixedVecs s sh' (Mixed (MapJust sh) a))
                     vecs)
+  mvecsUnsafeFreeze :: forall sh' s. IShX sh' -> MixedVecs s sh' (Shaped sh a) -> ST s (Mixed sh' (Shaped sh a))
+  mvecsUnsafeFreeze sh vecs =
+    coerce @(Mixed sh' (Mixed (MapJust sh) a))
+           @(Mixed sh' (Shaped sh a))
+      <$> mvecsUnsafeFreeze sh
+            (coerce @(MixedVecs s sh' (Shaped sh a))
+                    @(MixedVecs s sh' (Mixed (MapJust sh) a))
+                    vecs)
 
 instance (KnownShS sh, KnownElt a) => KnownElt (Shaped sh a) where
   memptyArrayUnsafe :: forall sh'. IShX sh' -> Mixed sh' (Shaped sh a)
@@ -181,6 +195,10 @@ instance (KnownShS sh, KnownElt a) => KnownElt (Shaped sh a) where
     | Dict <- lemKnownMapJust (Proxy @sh)
     = MV_Shaped <$> mvecsUnsafeNew idx arr
 
+  mvecsReplicate idx (Shaped arr)
+    | Dict <- lemKnownMapJust (Proxy @sh)
+    = MV_Shaped <$> mvecsReplicate idx arr
+
   mvecsNewEmpty _
     | Dict <- lemKnownMapJust (Proxy @sh)
     = MV_Shaped <$> mvecsNewEmpty (Proxy @(Mixed (MapJust sh) a))
@@ -242,14 +260,6 @@ satan2Array :: (FloatElt a, PrimElt a) => Shaped sh a -> Shaped sh a -> Shaped s
 satan2Array = liftShaped2 matan2Array
 
 
+{-# INLINE sshape #-}
 sshape :: forall sh a. Elt a => Shaped sh a -> ShS sh
-sshape (Shaped arr) = shsFromShX (mshape arr)
-
--- Needed already here, but re-exported in Data.Array.Nested.Convert.
-shsFromShX :: forall sh i. ShX (MapJust sh) i -> ShS sh
-shsFromShX ZSX = castWith (subst1 (unsafeCoerceRefl :: '[] :~: sh)) ZSS
-shsFromShX (SKnown n@SNat :$% (idx :: ShX mjshT i)) =
-  castWith (subst1 (sym (lemMapJustCons Refl))) $
-    n :$$ shsFromShX @(Tail sh) (castWith (subst2 (unsafeCoerceRefl :: mjshT :~: MapJust (Tail sh)))
-                                   idx)
-shsFromShX (SUnknown _ :$% _) = error "impossible"
+sshape (Shaped arr) = coerce (mshape arr)