3 files changed, 21 insertions, 12 deletions
diff --git a/src/Data/Array/Nested/Convert.hs b/src/Data/Array/Nested/Convert.hs
index c27c4c2..6be0f74 100644
--- a/src/Data/Array/Nested/Convert.hs
+++ b/src/Data/Array/Nested/Convert.hs
@@ -67,9 +67,12 @@ ixrFromIxS
     Prelude.. coerce
 
 ixrFromIxS' :: forall sh i. SNat (Rank sh) -> IxS sh i -> IxR (Rank sh) i
-ixrFromIxS' _ = coerce
-                Prelude.. (unsafeCoerce :: IxX (MapJust sh) i -> IxX (Replicate (Rank sh) Nothing) i)
-                Prelude.. coerce
+ixrFromIxS' _
+  | Refl <- lemRankReplicate (Proxy @(Rank sh))
+  , Refl <- lemRankMapJust (Proxy @sh)
+  = coerce
+    Prelude.. (coerceEqualRankListX :: ListX (MapJust sh) i -> ListX (Replicate (Rank sh) Nothing) i)
+    Prelude.. coerce
 
 -- ixrFromIxX re-exported
 
@@ -99,9 +102,12 @@ ixsFromIxR
     Prelude.. coerce
 
 ixsFromIxR' :: forall sh i. ShS sh -> IxR (Rank sh) i -> IxS sh i
-ixsFromIxR' _ = coerce
-                Prelude.. (unsafeCoerce :: IxX (Replicate (Rank sh) Nothing) i -> IxX (MapJust sh) i)
-                Prelude.. coerce
+ixsFromIxR' _
+  | Refl <- lemRankReplicate (Proxy @(Rank sh))
+  , Refl <- lemRankMapJust (Proxy @sh)
+  = coerce
+    Prelude.. (coerceEqualRankListX :: ListX (Replicate (Rank sh) Nothing) i -> ListX (MapJust sh) i)
+    Prelude.. coerce
 
 -- ixsFromIxX re-exported
 
diff --git a/src/Data/Array/Nested/Lemmas.hs b/src/Data/Array/Nested/Lemmas.hs
index 43cab3e..850f8ea 100644
--- a/src/Data/Array/Nested/Lemmas.hs
+++ b/src/Data/Array/Nested/Lemmas.hs
@@ -6,7 +6,11 @@
 {-# LANGUAGE TypeOperators #-}
 {-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
 {-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
-module Data.Array.Nested.Lemmas where
+module Data.Array.Nested.Lemmas (
+  module Data.Array.Nested.Lemmas,
+  lemReplicateSucc, lemMapJustEmpty, lemMapJustCons, lemMapJustHead,
+  lemRankMapJust
+) where
 
 import Data.Proxy
 import Data.Type.Equality
@@ -146,9 +150,6 @@ lemRankAppComm _ _ = unsafeCoerceRefl
 lemRankReplicate :: proxy n -> Rank (Replicate n (Nothing @Nat)) :~: n
 lemRankReplicate _ = unsafeCoerceRefl
 
-lemRankMapJust :: proxy sh -> Rank (MapJust sh) :~: Rank sh
-lemRankMapJust _ = unsafeCoerceRefl
-
 -- ** Related to MapJust and/or Permutation
 
 lemTakeLenMapJust :: Perm is -> ShS sh -> TakeLen is (MapJust sh) :~: MapJust (TakeLen is sh)
diff --git a/src/Data/Array/Nested/Shaped/Shape.hs b/src/Data/Array/Nested/Shaped/Shape.hs
index 3d4bf31..d1a72ea 100644
--- a/src/Data/Array/Nested/Shaped/Shape.hs
+++ b/src/Data/Array/Nested/Shaped/Shape.hs
@@ -314,10 +314,12 @@ shsLength :: ShS sh -> Int
 shsLength (ShS shx) = shxLength shx
 
 shsRank :: forall sh. ShS sh -> SNat (Rank sh)
-shsRank (ShS shx) =
-  gcastWith (unsafeCoerceRefl :: Rank (MapJust sh) :~: Rank sh) $
+shsRank (ShS shx) | Refl <- lemRankMapJust (Proxy @sh) =
   shxRank shx
 
+lemRankMapJust :: proxy sh -> Rank (MapJust sh) :~: Rank sh
+lemRankMapJust _ = unsafeCoerceRefl
+
 shsSize :: ShS sh -> Int
 shsSize (ShS sh) = shxSize sh