From 367179c65d0dfb942cf5e36bde2e0c062c5dad6c Mon Sep 17 00:00:00 2001
From: Tom Smeding <tom@tomsmeding.com>
Date: Mon, 7 Oct 2024 10:17:48 +0200
Subject: TestEquality for StaticShX

---
 src/Data/Array/Mixed/Shape.hs | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

(limited to 'src')

diff --git a/src/Data/Array/Mixed/Shape.hs b/src/Data/Array/Mixed/Shape.hs
index 1285aa1..a8f58ef 100644
--- a/src/Data/Array/Mixed/Shape.hs
+++ b/src/Data/Array/Mixed/Shape.hs
@@ -70,6 +70,14 @@ 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')
+  | Just Refl <- testEquality n m
+  , Just Refl <- listxEq sh sh'
+  = Just Refl
+listxEq _ _ = Nothing
+
 listxFmap :: (forall n. f n -> g n) -> ListX sh f -> ListX sh g
 listxFmap _ ZX = ZX
 listxFmap f (x ::% xs) = f x ::% listxFmap f xs
@@ -224,6 +232,11 @@ instance (NFData i, forall m. NFData (f m)) => NFData (SMayNat i f n) where
   rnf (SUnknown i) = rnf i
   rnf (SKnown x) = rnf x
 
+instance TestEquality f => TestEquality (SMayNat i f) where
+  testEquality SUnknown{} SUnknown{} = Just Refl
+  testEquality (SKnown n) (SKnown m) | Just Refl <- testEquality n m = Just Refl
+  testEquality _ _ = Nothing
+
 fromSMayNat :: (n ~ Nothing => i -> r)
             -> (forall m. n ~ Just m => f m -> r)
             -> SMayNat i f n -> r
@@ -388,6 +401,9 @@ infixr 3 :!%
 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
+
 ssxLength :: StaticShX sh -> Int
 ssxLength (StaticShX l) = listxLength l
 
-- 
cgit v1.2.3-70-g09d2