Implement ListX as [] with strict pattern synonyms

2 files changed, 130 insertions, 58 deletions

diff --git a/src/Data/Array/Nested/Mixed/ListX.hs b/src/Data/Array/Nested/Mixed/ListX.hs
new file mode 100644
index 0000000..e89d1c8
--- /dev/null
+++ b/src/Data/Array/Nested/Mixed/ListX.hs
@@ -0,0 +1,125 @@
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE CPP #-}
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE ImportQualifiedPost #-}
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE NoStarIsType #-}
+{-# LANGUAGE PatternSynonyms #-}
+{-# LANGUAGE PolyKinds #-}
+{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE RoleAnnotations #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE StandaloneKindSignatures #-}
+{-# LANGUAGE StrictData #-}
+{-# LANGUAGE TemplateHaskell #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE UnboxedTuples #-}
+{-# LANGUAGE UndecidableInstances #-}
+{-# 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
+
+import Control.DeepSeq (NFData(..))
+import Data.Foldable qualified as Foldable
+import Data.Kind (Type)
+import Data.Type.Equality
+import GHC.IsList (IsList)
+import GHC.IsList qualified as IsList
+import GHC.TypeLits
+#if !MIN_VERSION_GLASGOW_HASKELL(9,8,0,0)
+import GHC.TypeLits.Orphans ()
+#endif
+
+import Data.Array.Nested.Types
+
+-- * Mixed lists implementation
+
+type role ListX nominal representational
+type ListX :: [Maybe Nat] -> Type -> Type
+newtype ListX sh i = ListX [i]
+  -- data invariant: each element is in WHNF; the spine maybe be not forced
+  deriving (Eq, Ord, NFData, Foldable)
+
+pattern ZX :: forall sh i. () => sh ~ '[] => ListX sh i
+pattern ZX <- (listxNull -> Just Refl)
+  where ZX = ListX []
+
+{-# INLINE listxNull #-}
+listxNull :: ListX sh i -> Maybe (sh :~: '[])
+listxNull (ListX []) = Just unsafeCoerceRefl
+listxNull (ListX (_ : _)) = Nothing
+
+{-# INLINE (::%) #-}
+pattern (::%)
+  :: forall {sh1} {i}.
+     forall n sh. (n : sh ~ sh1)
+  => i -> ListX sh i -> ListX sh1 i
+pattern i ::% sh <- (listxUncons -> Just (UnconsListXRes sh i))
+  where !i ::% ListX !sh = ListX (i : sh)
+infixr 3 ::%
+
+data UnconsListXRes i sh1 =
+  forall n sh. (n : sh ~ sh1) => UnconsListXRes (ListX sh i) i
+{-# INLINE listxUncons #-}
+listxUncons :: forall sh1 i. ListX sh1 i -> Maybe (UnconsListXRes i sh1)
+listxUncons (ListX (i : sh')) = gcastWith (unsafeCoerceRefl :: Head sh1 ': Tail sh1 :~: sh1) $
+                                Just (UnconsListXRes (ListX @(Tail sh1) sh') i)
+listxUncons (ListX []) = Nothing
+
+{-# COMPLETE ZX, (::%) #-}
+
+{-# INLINE lazily #-}
+lazily :: ([a] -> [b]) -> ListX sh a -> ListX sh b
+lazily f (ListX l) = ListX $ f l
+
+{-# INLINE lazilyConcat #-}
+lazilyConcat :: ([a] -> [b] -> [c]) -> ListX sh a -> ListX sh' b -> ListX (sh ++ sh') c
+lazilyConcat f (ListX l) (ListX k) = ListX $ f l k
+
+{-# INLINE lazilyForce #-}
+lazilyForce :: ([a] -> [b]) -> ListX sh a -> ListX sh b
+lazilyForce f (ListX l) = let res = f l
+                          in foldr seq () res `seq` ListX res
+
+#ifdef OXAR_DEFAULT_SHOW_INSTANCES
+deriving instance Show i => Show (ListX sh i)
+#else
+instance Show i => Show (ListX sh i) where
+  showsPrec _ = listxShow shows
+#endif
+
+{-# INLINE listxShow #-}
+listxShow :: forall sh i. (i -> ShowS) -> ListX sh i -> ShowS
+listxShow f l = showString "[" . go "" l . showString "]"
+  where
+    go :: String -> ListX sh' i -> ShowS
+    go _ ZX = id
+    go prefix (x ::% xs) = showString prefix . f x . go "," xs
+
+-- This can't be derived, becauses the list needs to be fully evaluated,
+-- per data invariant. This version is faster than versions defined using
+-- (::%) or lazilyForce.
+instance Functor (ListX l) where
+  {-# INLINE fmap #-}
+  fmap f (ListX l) =
+    let fmap' [] = []
+        fmap' (x : xs) = let y = f x
+                             rest = fmap' xs
+                         in y `seq` rest `seq` y : rest
+    in ListX $ fmap' l
+
+-- | Very untyped: not even length is checked (at runtime).
+instance IsList (ListX sh i) where
+  type Item (ListX sh i) = i
+  {-# INLINE fromList #-}
+  fromList l = foldr seq () l `seq` ListX l
+  {-# INLINE toList #-}
+  toList = Foldable.toList
diff --git a/src/Data/Array/Nested/Mixed/Shape.hs b/src/Data/Array/Nested/Mixed/Shape.hs
index 2dfcc8c..611ec19 100644
--- a/src/Data/Array/Nested/Mixed/Shape.hs
+++ b/src/Data/Array/Nested/Mixed/Shape.hs
@@ -28,6 +28,7 @@
 module Data.Array.Nested.Mixed.Shape where
 
 import Control.DeepSeq (NFData(..))
+import Control.Exception (assert)
 import Data.Bifunctor (first)
 import Data.Coerce
 import Data.Foldable qualified as Foldable
@@ -43,6 +44,7 @@ import GHC.TypeLits
 import GHC.TypeLits.Orphans ()
 #endif
 
+import Data.Array.Nested.Mixed.ListX
 import Data.Array.Nested.Types
 
 
@@ -55,63 +57,14 @@ type family Rank sh where
 
 -- * Mixed lists
 
-type role ListX nominal representational
-type ListX :: [Maybe Nat] -> Type -> Type
-data ListX sh i where
-  ZX :: ListX '[] i
-  (::%) :: forall n sh {i}. i -> ListX sh i -> ListX (n : sh) i
-deriving instance Eq i => Eq (ListX sh i)
-deriving instance Ord i => Ord (ListX sh i)
-infixr 3 ::%
-
-#ifdef OXAR_DEFAULT_SHOW_INSTANCES
-deriving instance Show i => Show (ListX sh i)
-#else
-instance Show i => Show (ListX sh i) where
-  showsPrec _ = listxShow shows
-#endif
-
-instance NFData i => NFData (ListX sh i) where
-  rnf ZX = ()
-  rnf (x ::% l) = rnf x `seq` rnf l
-
-instance Functor (ListX l) where
-  {-# INLINE fmap #-}
-  fmap _ ZX = ZX
-  fmap f (x ::% xs) = f x ::% fmap f xs
-
-instance Foldable (ListX l) where
-  {-# INLINE foldMap #-}
-  foldMap _ ZX = mempty
-  foldMap f (x ::% xs) = f x <> foldMap f xs
-  {-# INLINE foldr #-}
-  foldr _ z ZX = z
-  foldr f z (x ::% xs) = f x (foldr f z xs)
-  null ZX = False
-  null _ = True
+{-# INLINE listxFromList #-}
+listxFromList :: StaticShX sh -> [i] -> ListX sh i
+listxFromList sh l = assert (ssxLength sh == length l) $ IsList.fromList l
 
 listxRank :: ListX sh i -> SNat (Rank sh)
 listxRank ZX = SNat
 listxRank (_ ::% l) | SNat <- listxRank l = SNat
 
-{-# INLINE listxShow #-}
-listxShow :: forall sh i. (i -> ShowS) -> ListX sh i -> ShowS
-listxShow f l = showString "[" . go "" l . showString "]"
-  where
-    go :: String -> ListX sh' i -> ShowS
-    go _ ZX = id
-    go prefix (x ::% xs) = showString prefix . f x . go "," xs
-
-listxFromList :: StaticShX sh -> [i] -> ListX sh i
-listxFromList topssh topl = go topssh topl
-  where
-    go :: StaticShX sh' -> [i] -> ListX sh' i
-    go ZKX [] = ZX
-    go (_ :!% sh) (i : is) = i ::% go sh is
-    go _ _ = error $ "listxFromList: Mismatched list length (type says "
-                       ++ show (ssxLength topssh) ++ ", list has length "
-                       ++ show (length topl) ++ ")"
-
 listxHead :: ListX (mn ': sh) i -> i
 listxHead (i ::% _) = i
 
@@ -772,12 +725,6 @@ shxFlatten = go (SNat @1)
     goUnknown acc (SKnown sn :$% sh) = goUnknown (acc * fromSNat' sn) sh
 
 
--- | Very untyped: only length is checked (at runtime).
-instance KnownShX sh => IsList (ListX sh i) where
-  type Item (ListX sh i) = i
-  fromList = listxFromList (knownShX @sh)
-  toList = Foldable.toList
-
 -- | 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