1 files changed, 139 insertions, 0 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..2c8a9cc
--- /dev/null
+++ b/src/Data/Array/Nested/Mixed/ListX.hs
@@ -0,0 +1,139 @@
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE CPP #-}
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE ImportQualifiedPost #-}
+{-# LANGUAGE NoStarIsType #-}
+{-# LANGUAGE PatternSynonyms #-}
+{-# LANGUAGE PolyKinds #-}
+{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE RoleAnnotations #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE StandaloneKindSignatures #-}
+{-# LANGUAGE StrictData #-}
+{-# 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, Rank, coerceEqualRankListX) 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
+
+
+-- | 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 on the list is in WHNF (the spine may be lazy)
+-- and the length of the list is the same as of the type-level shape.
+type role ListX nominal representational
+type ListX :: [Maybe Nat] -> Type -> Type
+newtype ListX sh i = ListX [i]
+  deriving (Eq, Ord, NFData, Foldable)
+
+{-# INLINE ZX #-}
+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 ::% l <- (listxUncons -> Just (UnconsListXRes i l))
+  where !i ::% ListX !l = ListX (i : l)
+infixr 3 ::%
+
+data UnconsListXRes i sh1 =
+  forall n sh. (n : sh ~ sh1) => UnconsListXRes i (ListX sh i)
+{-# INLINE listxUncons #-}
+listxUncons :: forall sh1 i. ListX sh1 i -> Maybe (UnconsListXRes i sh1)
+listxUncons (ListX (i : l)) = gcastWith (unsafeCoerceRefl :: Head sh1 ': Tail sh1 :~: sh1) $
+                              Just (UnconsListXRes i (ListX @(Tail sh1) l))
+listxUncons (ListX []) = Nothing
+
+{-# COMPLETE ZX, (::%) #-}
+
+-- | This function makes no attempt to stop you from breaking the data invariant for 'ListX'. If you do so, you must later ensure that the invariant is reinstated, for example using @'lazilyForce' 'id'@.
+{-# INLINE lazily #-}
+lazily :: ([a] -> [b]) -> ListX sh a -> ListX sh b
+lazily f (ListX l) = ListX $ f l
+
+-- | This function makes no attempt to stop you from breaking the data invariant for 'ListX'. If you do so, you must later ensure that the invariant is reinstated, for example using @'lazilyForce' 'id'@.
+{-# 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
+
+-- | This operation forces all elements of the @[b]@ list to restore the strictness part of the data invariant for 'ListX'. Note that ensuring the list has the right length is still the user's responsibility.
+{-# 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
+
+coerceEqualRankListX :: Rank sh ~ Rank sh' => ListX sh i -> ListX sh' i
+coerceEqualRankListX (ListX l) = ListX l