diff options
Diffstat (limited to 'src/Data/Array/Nested')
| -rw-r--r-- | src/Data/Array/Nested/Mixed/ListX.hs | 125 | ||||
| -rw-r--r-- | src/Data/Array/Nested/Mixed/Shape.hs | 63 | ||||
| -rw-r--r-- | src/Data/Array/Nested/Permutation.hs | 1 | ||||
| -rw-r--r-- | src/Data/Array/Nested/Ranked/Shape.hs | 13 | ||||
| -rw-r--r-- | src/Data/Array/Nested/Shaped/Shape.hs | 25 |
5 files changed, 144 insertions, 83 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 diff --git a/src/Data/Array/Nested/Permutation.hs b/src/Data/Array/Nested/Permutation.hs index 9254728..b6e5f47 100644 --- a/src/Data/Array/Nested/Permutation.hs +++ b/src/Data/Array/Nested/Permutation.hs @@ -31,6 +31,7 @@ import GHC.TypeNats qualified as TN import Data.Array.Nested.Mixed.Shape import Data.Array.Nested.Types +import Data.Array.Nested.Mixed.ListX -- * Permutations diff --git a/src/Data/Array/Nested/Ranked/Shape.hs b/src/Data/Array/Nested/Ranked/Shape.hs index 2166123..4eef090 100644 --- a/src/Data/Array/Nested/Ranked/Shape.hs +++ b/src/Data/Array/Nested/Ranked/Shape.hs @@ -28,6 +28,7 @@ module Data.Array.Nested.Ranked.Shape where import Control.DeepSeq (NFData(..)) +import Control.Exception (assert) import Data.Coerce (coerce) import Data.Foldable qualified as Foldable import Data.Kind (Type) @@ -40,6 +41,7 @@ import GHC.TypeLits import GHC.TypeNats qualified as TN import Data.Array.Nested.Lemmas +import Data.Array.Nested.Mixed.ListX import Data.Array.Nested.Mixed.Shape import Data.Array.Nested.Permutation import Data.Array.Nested.Types @@ -121,15 +123,10 @@ listrAppend :: forall n m i. ListR n i -> ListR m i -> ListR (n + m) i listrAppend ZR sh = sh listrAppend (x ::: xs) sh = x ::: listrAppend xs sh +{-# INLINE listrFromList #-} listrFromList :: SNat n -> [i] -> ListR n i -listrFromList topsn topl = go topsn topl - where - go :: SNat n' -> [i] -> ListR n' i - go SZ [] = ZR - go (SS n) (i : is) = i ::: go n is - go _ _ = error $ "listrFromList: Mismatched list length (type says " - ++ show (fromSNat topsn) ++ ", list has length " - ++ show (length topl) ++ ")" +listrFromList topsn topl = assert (fromSNat' topsn == length topl) + $ ListR $ IsList.fromList topl listrHead :: ListR (n + 1) i -> i listrHead (i ::: _) = i diff --git a/src/Data/Array/Nested/Shaped/Shape.hs b/src/Data/Array/Nested/Shaped/Shape.hs index 274f954..13596a7 100644 --- a/src/Data/Array/Nested/Shaped/Shape.hs +++ b/src/Data/Array/Nested/Shaped/Shape.hs @@ -26,6 +26,7 @@ module Data.Array.Nested.Shaped.Shape where import Control.DeepSeq (NFData(..)) +import Control.Exception (assert) import Data.Array.Shape qualified as O import Data.Coerce (coerce) import Data.Foldable qualified as Foldable @@ -37,6 +38,7 @@ import GHC.IsList (IsList) import GHC.IsList qualified as IsList import GHC.TypeLits +import Data.Array.Nested.Mixed.ListX import Data.Array.Nested.Mixed.Shape import Data.Array.Nested.Permutation import Data.Array.Nested.Types @@ -93,26 +95,15 @@ listsRank :: ListS sh i -> SNat (Rank sh) listsRank ZS = SNat listsRank (_ ::$ sh) = snatSucc (listsRank sh) +{-# INLINE listsFromList #-} listsFromList :: ShS sh -> [i] -> ListS sh i -listsFromList topsh topl = go topsh topl - where - go :: ShS sh' -> [i] -> ListS sh' i - go ZSS [] = ZS - go (_ :$$ sh) (i : is) = i ::$ go sh is - go _ _ = error $ "listsFromList: Mismatched list length (type says " - ++ show (shsLength topsh) ++ ", list has length " - ++ show (length topl) ++ ")" +listsFromList sh l = assert (shsLength sh == length l) + $ ListS $ IsList.fromList l -{-# INLINEABLE listsFromListS #-} +{-# INLINE listsFromListS #-} listsFromListS :: ListS sh i0 -> [i] -> ListS sh i -listsFromListS topl0 topl = go topl0 topl - where - go :: ListS sh i0 -> [i] -> ListS sh i - go ZS [] = ZS - go (_ ::$ l0) (i : is) = i ::$ go l0 is - go _ _ = error $ "listsFromListS: Mismatched list length (the model says " - ++ show (length topl0) ++ ", list has length " - ++ show (length topl) ++ ")" +listsFromListS sh l = assert (length sh == length l) + $ ListS $ IsList.fromList l listsHead :: ListS (n : sh) i -> i listsHead (i ::$ _) = i |