diff options
Diffstat (limited to 'src/Data/Array/Nested/Mixed.hs')
| -rw-r--r-- | src/Data/Array/Nested/Mixed.hs | 45 |
1 files changed, 28 insertions, 17 deletions
diff --git a/src/Data/Array/Nested/Mixed.hs b/src/Data/Array/Nested/Mixed.hs index 221393f..250c999 100644 --- a/src/Data/Array/Nested/Mixed.hs +++ b/src/Data/Array/Nested/Mixed.hs @@ -16,6 +16,7 @@ {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE TypeOperators #-} {-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE UndecidableSuperClasses #-} {-# LANGUAGE ViewPatterns #-} module Data.Array.Nested.Mixed where @@ -29,7 +30,7 @@ import Data.Bifunctor (bimap) import Data.Coerce import Data.Foldable (toList) import Data.Int -import Data.Kind (Type) +import Data.Kind (Type, Constraint) import Data.List.NonEmpty (NonEmpty(..)) import Data.List.NonEmpty qualified as NE import Data.Proxy @@ -290,7 +291,9 @@ matan2Array = mliftNumElt2 (liftO2 . floatEltAtan2) -- | Allowable element types in a mixed array, and by extension in a 'Ranked' or -- 'Shaped' array. Note the polymorphic instance for 'Elt' of @'Primitive' -- a@; see the documentation for 'Primitive' for more details. -class Elt a where +class EltC a => Elt a where + type EltC a :: Constraint + -- ====== PUBLIC METHODS ====== -- mshape :: Mixed sh a -> IShX sh @@ -383,7 +386,9 @@ class Elt a where -- This class is (currently) only required for 'mgenerate', -- 'Data.Array.Nested.Ranked.rgenerate' and -- 'Data.Array.Nested.Shaped.sgenerate'. -class Elt a => KnownElt a where +class (Elt a, KnownEltC a) => KnownElt a where + type KnownEltC a :: Constraint + -- | Create an empty array. The given shape must have size zero; this may or may not be checked. memptyArrayUnsafe :: IShX sh -> Mixed sh a @@ -396,6 +401,8 @@ class Elt a => KnownElt a where -- Arrays of scalars are basically just arrays of scalars. instance Storable a => Elt (Primitive a) where + type EltC (Primitive a) = Storable a + mshape (M_Primitive sh _) = sh mindex (M_Primitive _ a) i = Primitive (X.index a i) mindexPartial (M_Primitive sh a) i = M_Primitive (shxDropIx sh i) (X.indexPartial a i) @@ -484,6 +491,7 @@ deriving via Primitive Float instance Elt Float deriving via Primitive () instance Elt () instance Storable a => KnownElt (Primitive a) where + type KnownEltC (Primitive a) = () memptyArrayUnsafe sh = M_Primitive sh (X.empty sh) mvecsUnsafeNew sh _ = MV_Primitive <$> VSM.unsafeNew (shxSize sh) mvecsNewEmpty _ = MV_Primitive <$> VSM.unsafeNew 0 @@ -500,6 +508,8 @@ deriving via Primitive () instance KnownElt () -- Arrays of pairs are pairs of arrays. instance (Elt a, Elt b) => Elt (a, b) where + type EltC (a, b) = (Elt a, Elt b) + mshape (M_Tup2 a _) = mshape a mindex (M_Tup2 a b) i = (mindex a i, mindex b i) mindexPartial (M_Tup2 a b) i = M_Tup2 (mindexPartial a i) (mindexPartial b i) @@ -543,12 +553,15 @@ instance (Elt a, Elt b) => Elt (a, b) where mvecsFreeze sh (MV_Tup2 a b) = M_Tup2 <$> mvecsFreeze sh a <*> mvecsFreeze sh b instance (KnownElt a, KnownElt b) => KnownElt (a, b) where + type KnownEltC (a, b) = (KnownEltC a, KnownEltC b) memptyArrayUnsafe sh = M_Tup2 (memptyArrayUnsafe sh) (memptyArrayUnsafe sh) mvecsUnsafeNew sh (x, y) = MV_Tup2 <$> mvecsUnsafeNew sh x <*> mvecsUnsafeNew sh y mvecsNewEmpty _ = MV_Tup2 <$> mvecsNewEmpty (Proxy @a) <*> mvecsNewEmpty (Proxy @b) -- Arrays of arrays are just arrays, but with more dimensions. instance Elt a => Elt (Mixed sh' a) where + type EltC (Mixed sh' a) = Elt a + -- TODO: this is quadratic in the nesting depth because it repeatedly -- truncates the shape vector to one a little shorter. Fix with a -- moverlongShape method, a prefix of which is mshape. @@ -681,6 +694,7 @@ instance Elt a => Elt (Mixed sh' a) where mvecsFreeze sh (MV_Nest sh' vecs) = M_Nest sh <$> mvecsFreeze (shxAppend sh sh') vecs instance (KnownShX sh', KnownElt a) => KnownElt (Mixed sh' a) where + type KnownEltC (Mixed sh' a) = (KnownShX sh', KnownElt a) memptyArrayUnsafe sh = M_Nest sh (memptyArrayUnsafe (shxAppend sh (shxCompleteZeros (knownShX @sh')))) mvecsUnsafeNew sh example @@ -784,14 +798,10 @@ mtoVector arr = mtoVectorP (toPrimitive arr) mfromList1 :: Elt a => NonEmpty a -> Mixed '[Nothing] a mfromList1 = mfromListOuter . fmap mscalar -- TODO: optimise? -mfromList1Prim :: PrimElt a => [a] -> Mixed '[Nothing] a -mfromList1Prim l = - let ssh = SUnknown () :!% ZKX - xarr = X.fromList1 ssh l - in fromPrimitive $ M_Primitive (X.shape ssh xarr) xarr - -mtoList1 :: Elt a => Mixed '[n] a -> [a] -mtoList1 = map munScalar . mtoListOuter +-- This forall is there so that a simple type application can constrain the +-- shape, in case the user wants to use OverloadedLists for the shape. +mfromListLinear :: forall sh a. Elt a => IShX sh -> NonEmpty a -> Mixed sh a +mfromListLinear sh l = mreshape sh (mfromList1 l) mfromListPrim :: PrimElt a => [a] -> Mixed '[Nothing] a mfromListPrim l = @@ -804,10 +814,8 @@ mfromListPrimLinear sh l = let M_Primitive _ xarr = toPrimitive (mfromListPrim l) in fromPrimitive $ M_Primitive sh (X.reshape (SUnknown () :!% ZKX) sh xarr) --- This forall is there so that a simple type application can constrain the --- shape, in case the user wants to use OverloadedLists for the shape. -mfromListLinear :: forall sh a. Elt a => IShX sh -> NonEmpty a -> Mixed sh a -mfromListLinear sh l = mreshape sh (mfromList1 l) +mtoList :: Elt a => Mixed '[n] a -> [a] +mtoList = map munScalar . mtoListOuter mtoListLinear :: Elt a => Mixed sh a -> [a] mtoListLinear arr = map (mindex arr) (shxEnum (mshape arr)) -- TODO: optimise @@ -821,8 +829,11 @@ mnest ssh arr = M_Nest (fst (shxSplitApp (Proxy @sh') ssh (mshape arr))) arr munNest :: Mixed sh (Mixed sh' a) -> Mixed (sh ++ sh') a munNest (M_Nest _ arr) = arr -mzip :: Mixed sh a -> Mixed sh b -> Mixed sh (a, b) -mzip = M_Tup2 +-- | The arguments must have equal shapes. If they do not, an error is raised. +mzip :: (Elt a, Elt b) => Mixed sh a -> Mixed sh b -> Mixed sh (a, b) +mzip a b + | Just Refl <- shxEqual (mshape a) (mshape b) = M_Tup2 a b + | otherwise = error "mzip: unequal shapes" munzip :: Mixed sh (a, b) -> (Mixed sh a, Mixed sh b) munzip (M_Tup2 a b) = (a, b) |