From 390af124e6cb50c4a2cd9006662fc26eef02889a Mon Sep 17 00:00:00 2001
From: Tom Smeding <tom@tomsmeding.com>
Date: Sun, 19 May 2024 12:55:58 +0200
Subject: Some IsList instances

---
 src/Data/Array/Mixed.hs | 48 +++++++++++++++++++++++++++++++++---------------
 1 file changed, 33 insertions(+), 15 deletions(-)

(limited to 'src/Data/Array/Mixed.hs')

diff --git a/src/Data/Array/Mixed.hs b/src/Data/Array/Mixed.hs
index 31a4e69..398a3de 100644
--- a/src/Data/Array/Mixed.hs
+++ b/src/Data/Array/Mixed.hs
@@ -30,6 +30,7 @@ import Data.Bifunctor (first)
 import Data.Coerce
 import Data.Functor.Const
 import Data.Kind
+import Data.Monoid (Sum(..))
 import Data.Proxy
 import Data.Type.Bool
 import Data.Type.Equality
@@ -114,6 +115,9 @@ foldListX :: Monoid m => (forall n. f n -> m) -> ListX sh f -> m
 foldListX _ ZX = mempty
 foldListX f (x ::% xs) = f x <> foldListX f xs
 
+lengthListX :: ListX sh f -> Int
+lengthListX = getSum . foldListX (\_ -> Sum 1)
+
 showListX :: forall sh f. (forall n. f n -> ShowS) -> ListX sh f -> ShowS
 showListX f l = showString "[" . go "" l . showString "]"
   where
@@ -193,8 +197,7 @@ instance Functor (ShX sh) where
   fmap f (ShX l) = ShX (fmapListX (fromSMayNat (SUnknown . f) SKnown) l)
 
 lengthShX :: ShX sh i -> Int
-lengthShX ZSX = 0
-lengthShX (_ :$% sh) = 1 + lengthShX sh
+lengthShX (ShX l) = lengthListX l
 
 
 -- | The part of a shape that is statically known.
@@ -218,6 +221,9 @@ infixr 3 :!%
 instance Show (StaticShX sh) where
   showsPrec _ (StaticShX l) = showListX (fromSMayNat shows (shows . fromSNat)) l
 
+lengthStaticShX :: StaticShX sh -> Int
+lengthStaticShX (StaticShX l) = lengthListX l
+
 
 -- | Evidence for the static part of a shape. This pops up only when you are
 -- polymorphic in the element type of an array.
@@ -228,36 +234,48 @@ instance (KnownNat n, KnownShX sh) => KnownShX (Just n : sh) where knownShX = SK
 instance KnownShX sh => KnownShX (Nothing : sh) where knownShX = SUnknown () :!% knownShX
 
 
--- | Very untyped; length is checked at runtime.
+-- | Very untyped: only length is checked (at runtime).
 instance KnownShX sh => IsList (ListX sh (Const i)) where
   type Item (ListX sh (Const i)) = i
-  fromList = go (knownShX @sh)
+  fromList topl = go (knownShX @sh) topl
     where
       go :: StaticShX sh' -> [i] -> ListX sh' (Const i)
       go ZKX [] = ZX
       go (_ :!% sh) (i : is) = Const i ::% go sh is
-      go _ _ = error "IsList(ListX): Mismatched list length"
+      go _ _ = error $ "IsList(ListX): Mismatched list length (type says "
+                         ++ show typelen ++ ", list has length "
+                         ++ show (length topl) ++ ")"
+        where typelen = let StaticShX l = knownShX @sh in lengthListX l
   toList = go
     where
       go :: ListX sh' (Const i) -> [i]
       go ZX = []
       go (Const i ::% is) = i : go is
 
--- | Very untyped; length is checked at runtime, and index bounds are *not checked*.
+-- | Very untyped: only length is checked (at runtime), index bounds are __not checked__.
 instance KnownShX sh => IsList (IxX sh i) where
   type Item (IxX sh i) = i
   fromList = IxX . fromList
   toList (IxX l) = toList l
 
--- | Very untyped; length is checked at runtime, and known dimensions are *not checked*.
--- instance KnownShX sh => IsList (ShX sh i) where
---   type Item (ShX sh i) = i
---   fromList = ShX . fmapListX (\(Const i) -> _) . fromList
---   toList = go
---     where
---       go :: ShX sh' i -> [i]
---       go ZSX = []
---       go (Const i :$% is) = i : go is
+-- | Untyped: length and known dimensions are checked (at runtime).
+instance KnownShX sh => IsList (ShX sh Int) where
+  type Item (ShX sh Int) = Int
+  fromList = ShX . go (knownShX @sh)
+    where
+      go :: StaticShX sh' -> [Int] -> ListX sh' (SMayNat Int SNat)
+      go ZKX [] = ZX
+      go (SKnown sn :!% sh) (i : is)
+        | i == fromSNat' sn = SKnown sn ::% go sh is
+        | otherwise = error $ "IsList(ShX): Value does not match typing (type says "
+                                ++ show (fromSNat' sn) ++ ", list contains " ++ show i ++ ")"
+      go (SUnknown () :!% sh) (i : is) = SUnknown i ::% go sh is
+      go _ _ = error "IsList(ShX): Mismatched list length"
+  toList = go
+    where
+      go :: ShX sh' Int -> [Int]
+      go ZSX = []
+      go (smn :$% sh) = fromSMayNat' smn : go sh
 
 
 type family Rank sh where
-- 
cgit v1.2.3-70-g09d2