3 files changed, 14 insertions, 25 deletions
diff --git a/src/Data/Array/Nested/Convert.hs b/src/Data/Array/Nested/Convert.hs
index d4d1cea..32248c4 100644
--- a/src/Data/Array/Nested/Convert.hs
+++ b/src/Data/Array/Nested/Convert.hs
@@ -15,7 +15,7 @@
 module Data.Array.Nested.Convert (
   -- * Shape\/index\/list casting functions
   -- ** To ranked
-  ixrFromIxS, ixrFromIxX, shrFromShS, shrFromShX, shrFromShX2,
+  ixrFromIxS, ixrFromIxX, shrFromShS, shrFromShXAnyShape, shrFromShX,
   listrCast, ixrCast, shrCast,
   -- ** To shaped
   ixsFromIxR, ixsFromIxR', ixsFromIxX, ixsFromIxX', withShsFromShR, shsFromShX, withShsFromShX, shsFromSSX,
@@ -70,8 +70,13 @@ shrFromShS :: ShS sh -> IShR (Rank sh)
 shrFromShS ZSS = ZSR
 shrFromShS (n :$$ sh) = fromSNat' n :$: shrFromShS sh
 
--- shrFromShX re-exported
--- shrFromShX2 re-exported
+shrFromShXAnyShape :: IShX sh -> IShR (Rank sh)
+shrFromShXAnyShape ZSX = ZSR
+shrFromShXAnyShape (n :$% idx) = fromSMayNat' n :$: shrFromShXAnyShape idx
+
+shrFromShX :: IShX (Replicate n Nothing) -> IShR n
+shrFromShX = coerce
+
 -- listrCast re-exported
 -- ixrCast re-exported
 -- shrCast re-exported
@@ -146,10 +151,7 @@ ixxFromIxS :: IxS sh i -> IxX (MapJust sh) i
 ixxFromIxS = unsafeCoerce
 
 shxFromShR :: ShR n i -> ShX (Replicate n Nothing) i
-shxFromShR ZSR = ZSX
-shxFromShR (n :$: (idx :: ShR m i)) =
-  castWith (subst2 @ShX @i (lemReplicateSucc @(Nothing @Nat) (Proxy @m)))
-    (SUnknown n :$% shxFromShR idx)
+shxFromShR = coerce
 
 shxFromShS :: ShS sh -> IShX (MapJust sh)
 shxFromShS = coerce
diff --git a/src/Data/Array/Nested/Ranked.hs b/src/Data/Array/Nested/Ranked.hs
index d687983..ccbab63 100644
--- a/src/Data/Array/Nested/Ranked.hs
+++ b/src/Data/Array/Nested/Ranked.hs
@@ -220,7 +220,7 @@ rfromOrthotope sn arr
 
 rtoOrthotope :: forall a n. PrimElt a => Ranked n a -> S.Array n a
 rtoOrthotope (rtoPrimitive -> Ranked (M_Primitive sh (XArray arr)))
-  | Refl <- lemRankReplicate (shrRank $ shrFromShX2 @n sh)
+  | Refl <- lemRankReplicate (shrRank $ shrFromShX @n sh)
   = arr
 
 runScalar :: Elt a => Ranked 0 a -> a
@@ -345,10 +345,10 @@ rdot :: (PrimElt a, NumElt a) => Ranked n a -> Ranked n a -> a
 rdot = coerce mdot
 
 rtoXArrayPrimP :: Ranked n (Primitive a) -> (IShR n, XArray (Replicate n Nothing) a)
-rtoXArrayPrimP (Ranked arr) = first shrFromShX2 (mtoXArrayPrimP arr)
+rtoXArrayPrimP (Ranked arr) = first shrFromShX (mtoXArrayPrimP arr)
 
 rtoXArrayPrim :: PrimElt a => Ranked n a -> (IShR n, XArray (Replicate n Nothing) a)
-rtoXArrayPrim (Ranked arr) = first shrFromShX2 (mtoXArrayPrim arr)
+rtoXArrayPrim (Ranked arr) = first shrFromShX (mtoXArrayPrim arr)
 
 rfromXArrayPrimP :: SNat n -> XArray (Replicate n Nothing) a -> Ranked n (Primitive a)
 rfromXArrayPrimP sn arr = Ranked (mfromXArrayPrimP (ssxFromShX (X.shape (ssxFromSNat sn) arr)) arr)
diff --git a/src/Data/Array/Nested/Ranked/Base.hs b/src/Data/Array/Nested/Ranked/Base.hs
index 5c696f3..236cb05 100644
--- a/src/Data/Array/Nested/Ranked/Base.hs
+++ b/src/Data/Array/Nested/Ranked/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)
@@ -135,7 +134,7 @@ instance Elt a => Elt (Ranked n a) where
 
   type ShapeTree (Ranked n a) = (IShR n, ShapeTree a)
 
-  mshapeTree (Ranked arr) = first shrFromShX2 (mshapeTree arr)
+  mshapeTree (Ranked arr) = first coerce (mshapeTree arr)
 
   mshapeTreeEq _ (sh1, t1) (sh2, t2) = sh1 == sh2 && mshapeTreeEq (Proxy @a) t1 t2
 
@@ -259,19 +258,7 @@ ratan2Array = liftRanked2 matan2Array
 
 
 rshape :: Elt a => Ranked n a -> IShR n
-rshape (Ranked arr) = shrFromShX2 (mshape arr)
+rshape (Ranked arr) = coerce (mshape arr)
 
 rrank :: Elt a => Ranked n a -> SNat n
 rrank = shrRank . rshape
-
--- Needed already here, but re-exported in Data.Array.Nested.Convert.
-shrFromShX :: forall sh. IShX sh -> IShR (Rank sh)
-shrFromShX ZSX = ZSR
-shrFromShX (n :$% idx) = fromSMayNat' n :$: shrFromShX idx
-
--- Needed already here, but re-exported in Data.Array.Nested.Convert.
--- | Convenience wrapper around 'shrFromShX' that applies 'lemRankReplicate'.
-shrFromShX2 :: forall n. IShX (Replicate n Nothing) -> IShR n
-shrFromShX2 sh
-  | Refl <- lemRankReplicate (Proxy @n)
-  = shrFromShX sh