More consistent equality functions on shapes

1 files changed, 42 insertions, 21 deletions

diff --git a/src/Data/Array/Mixed/Shape.hs b/src/Data/Array/Mixed/Shape.hs
index e5f8b67..16f62fe 100644
--- a/src/Data/Array/Mixed/Shape.hs
+++ b/src/Data/Array/Mixed/Shape.hs
@@ -71,13 +71,28 @@ listxUncons :: ListX sh1 f -> Maybe (UnconsListXRes f sh1)
 listxUncons (i ::% shl') = Just (UnconsListXRes shl' i)
 listxUncons ZX = Nothing
 
-listxEq :: TestEquality f => ListX sh f -> ListX sh' f -> Maybe (sh :~: sh')
-listxEq ZX ZX = Just Refl
-listxEq (n ::% sh) (m ::% sh')
+-- | This checks only whether the types are equal; if the elements of the list
+-- are not singletons, their values may still differ. This corresponds to
+-- 'testEquality', except on the penultimate type parameter.
+listxEqType :: TestEquality f => ListX sh f -> ListX sh' f -> Maybe (sh :~: sh')
+listxEqType ZX ZX = Just Refl
+listxEqType (n ::% sh) (m ::% sh')
   | Just Refl <- testEquality n m
-  , Just Refl <- listxEq sh sh'
+  , Just Refl <- listxEqType sh sh'
   = Just Refl
-listxEq _ _ = Nothing
+listxEqType _ _ = Nothing
+
+-- | This checks whether the two lists actually contain equal values. This is
+-- more than 'testEquality', and corresponds to @geq@ from @Data.GADT.Compare@
+-- in the @some@ package (except on the penultimate type parameter).
+listxEqual :: (TestEquality f, forall n. Eq (f n)) => ListX sh f -> ListX sh' f -> Maybe (sh :~: sh')
+listxEqual ZX ZX = Just Refl
+listxEqual (n ::% sh) (m ::% sh')
+  | Just Refl <- testEquality n m
+  , n == m
+  , Just Refl <- listxEqual sh sh'
+  = Just Refl
+listxEqual _ _ = Nothing
 
 listxFmap :: (forall n. f n -> g n) -> ListX sh f -> ListX sh g
 listxFmap _ ZX = ZX
@@ -293,11 +308,26 @@ instance NFData i => NFData (ShX sh i) where
 shxLength :: ShX sh i -> Int
 shxLength (ShX l) = listxLength l
 
-shxRank :: ShX sh f -> SNat (Rank sh)
+shxRank :: ShX sh i -> SNat (Rank sh)
 shxRank (ShX list) = listxRank list
 
--- | This is more than @geq@: it also checks that the integers (the unknown
--- dimensions) are the same.
+-- | This checks only whether the types are equal; unknown dimensions might
+-- still differ. This corresponds to 'testEquality', except on the penultimate
+-- type parameter.
+shxEqType :: ShX sh i -> ShX sh' i -> Maybe (sh :~: sh')
+shxEqType ZSX ZSX = Just Refl
+shxEqType (SKnown n@SNat :$% sh) (SKnown m@SNat :$% sh')
+  | Just Refl <- sameNat n m
+  , Just Refl <- shxEqType sh sh'
+  = Just Refl
+shxEqType (SUnknown _ :$% sh) (SUnknown _ :$% sh')
+  | Just Refl <- shxEqType sh sh'
+  = Just Refl
+shxEqType _ _ = Nothing
+
+-- | This checks whether all dimensions have the same value. This is more than
+-- 'testEquality', and corresponds to @geq@ from @Data.GADT.Compare@ in the
+-- @some@ package (except on the penultimate type parameter).
 shxEqual :: Eq i => ShX sh i -> ShX sh' i -> Maybe (sh :~: sh')
 shxEqual ZSX ZSX = Just Refl
 shxEqual (SKnown n@SNat :$% sh) (SKnown m@SNat :$% sh')
@@ -417,23 +447,14 @@ instance Show (StaticShX sh) where
   showsPrec _ (StaticShX l) = listxShow (fromSMayNat shows (shows . fromSNat)) l
 
 instance TestEquality StaticShX where
-  testEquality (StaticShX l1) (StaticShX l2) = listxEq l1 l2
+  testEquality (StaticShX l1) (StaticShX l2) = listxEqType l1 l2
 
 ssxLength :: StaticShX sh -> Int
 ssxLength (StaticShX l) = listxLength l
 
--- | This suffices as an implementation of @geq@ in the @Data.GADT.Compare@
--- class of the @some@ package.
-ssxGeq :: StaticShX sh -> StaticShX sh' -> Maybe (sh :~: sh')
-ssxGeq ZKX ZKX = Just Refl
-ssxGeq (SKnown n@SNat :!% sh) (SKnown m@SNat :!% sh')
-  | Just Refl <- sameNat n m
-  , Just Refl <- ssxGeq sh sh'
-  = Just Refl
-ssxGeq (SUnknown () :!% sh) (SUnknown () :!% sh')
-  | Just Refl <- ssxGeq sh sh'
-  = Just Refl
-ssxGeq _ _ = Nothing
+-- | @ssxEqType = 'testEquality'@. Provided for consistency.
+ssxEqType :: StaticShX sh -> StaticShX sh' -> Maybe (sh :~: sh')
+ssxEqType = testEquality
 
 ssxAppend :: StaticShX sh -> StaticShX sh' -> StaticShX (sh ++ sh')
 ssxAppend ZKX sh' = sh'