Make equal-rank-coercibility part of the interface of ListXmove-Rank

4 files changed, 31 insertions, 18 deletions

diff --git a/src/Data/Array/Nested/Convert.hs b/src/Data/Array/Nested/Convert.hs
index 671f40f..4ee2a71 100644
--- a/src/Data/Array/Nested/Convert.hs
+++ b/src/Data/Array/Nested/Convert.hs
@@ -46,6 +46,7 @@ import Unsafe.Coerce (unsafeCoerce)
 
 import Data.Array.Nested.Lemmas
 import Data.Array.Nested.Mixed
+import Data.Array.Nested.Mixed.ListX
 import Data.Array.Nested.Mixed.Shape
 import Data.Array.Nested.Ranked.Base
 import Data.Array.Nested.Ranked.Shape
@@ -58,9 +59,12 @@ import Data.Array.Nested.Types
 -- * To ranked
 
 ixrFromIxS :: forall sh i. 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
 
@@ -82,9 +86,12 @@ shrFromShX = coerce
 -- * To shaped
 
 ixsFromIxR :: forall sh i. 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 e61b148..43cab3e 100644
--- a/src/Data/Array/Nested/Lemmas.hs
+++ b/src/Data/Array/Nested/Lemmas.hs
@@ -146,9 +146,8 @@ lemRankAppComm _ _ = unsafeCoerceRefl
 lemRankReplicate :: proxy n -> Rank (Replicate n (Nothing @Nat)) :~: n
 lemRankReplicate _ = unsafeCoerceRefl
 
-lemRankMapJust :: ShS sh -> Rank (MapJust sh) :~: Rank sh
-lemRankMapJust ZSS = Refl
-lemRankMapJust (_ :$$ sh') | Refl <- lemRankMapJust sh' = Refl
+lemRankMapJust :: proxy sh -> Rank (MapJust sh) :~: Rank sh
+lemRankMapJust _ = unsafeCoerceRefl
 
 -- ** Related to MapJust and/or Permutation
 
diff --git a/src/Data/Array/Nested/Mixed/ListX.hs b/src/Data/Array/Nested/Mixed/ListX.hs
index 51db341..e4e224b 100644
--- a/src/Data/Array/Nested/Mixed/ListX.hs
+++ b/src/Data/Array/Nested/Mixed/ListX.hs
@@ -23,7 +23,7 @@
 {-# LANGUAGE ViewPatterns #-}
 {-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
 {-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
-module Data.Array.Nested.Mixed.ListX (ListX, pattern ZX, pattern (::%), listxShow, lazily, lazilyConcat, lazilyForce) where
+module Data.Array.Nested.Mixed.ListX (ListX, pattern ZX, pattern (::%), listxShow, lazily, lazilyConcat, lazilyForce, Rank, coerceEqualRankListX) where
 
 import Control.DeepSeq (NFData(..))
 import Data.Foldable qualified as Foldable
@@ -38,6 +38,14 @@ import GHC.TypeLits.Orphans ()
 
 import Data.Array.Nested.Types
 
+
+-- | The length of a type-level list. If the argument is a shape, then the
+-- result is the rank of that shape.
+type family Rank sh where
+  Rank '[] = 0
+  Rank (_ : sh) = Rank sh + 1
+
+
 -- * Mixed lists implementation
 
 -- | Data invariant: each element is in WHNF (the spine may be be not forced).
@@ -132,3 +140,6 @@ instance IsList (ListX sh i) where
   fromList l = foldr seq () l `seq` ListX l
   {-# INLINE toList #-}
   toList = Foldable.toList
+
+coerceEqualRankListX :: Rank sh ~ Rank sh' => ListX sh i -> ListX sh' i
+coerceEqualRankListX (ListX l) = ListX l
diff --git a/src/Data/Array/Nested/Mixed/Shape.hs b/src/Data/Array/Nested/Mixed/Shape.hs
index a5e3ced..c678c17 100644
--- a/src/Data/Array/Nested/Mixed/Shape.hs
+++ b/src/Data/Array/Nested/Mixed/Shape.hs
@@ -24,7 +24,10 @@
 {-# LANGUAGE ViewPatterns #-}
 {-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise #-}
 {-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}
-module Data.Array.Nested.Mixed.Shape where
+module Data.Array.Nested.Mixed.Shape (
+  module Data.Array.Nested.Mixed.Shape,
+  Rank,
+) where
 
 import Control.DeepSeq (NFData(..))
 import Control.Exception (assert)
@@ -47,13 +50,6 @@ import Data.Array.Nested.Mixed.ListX
 import Data.Array.Nested.Types
 
 
--- | The length of a type-level list. If the argument is a shape, then the
--- result is the rank of that shape.
-type family Rank sh where
-  Rank '[] = 0
-  Rank (_ : sh) = Rank sh + 1
-
-
 -- * Mixed lists
 
 {-# INLINE listxFromList #-}